Coverage for src/ensae_teaching_cs/helpers/matplotlib_helper_xyz.py: 86%
80 statements
« prev ^ index » next coverage.py v7.1.0, created at 2023-04-28 06:23 +0200
« prev ^ index » next coverage.py v7.1.0, created at 2023-04-28 06:23 +0200
1"""
2@file
3@brief scatter plots
4"""
5import numpy
6from matplotlib.tri import Triangulation
7from matplotlib.pyplot import Normalize
8from .colorsdef import colors_definition
11def scatter_xy_id(xy_id, legend=None, ax=None, **options):
12 """
13 Creates a scatter plot with a different color for each zone id.
14 The function requires :epkg:`matplotlib`.
16 @param xy_id list of 3-uple *(x, y, zone_id)*
17 @param legend dictionary {id: legend } or None if there is not any
18 @param ax existing graph to plot on (can be None)
19 @param options others options: xlabel, ylabel, title, marker, figsize (if ax is None)
20 @return fig, ax (fig is None if ax was sent to the function)
22 .. plot::
23 :include-source:
25 import random
26 def generate_gauss(x, y, sigma, N=1000):
27 res = []
28 for i in range(N):
29 u = random.gauss(0, 1)
30 a = sigma * u + x
31 b = sigma * random.gauss(0, 1) + y + u
32 res.append((a, b))
33 return res
34 nuage1 = generate_gauss(0, 0, 3)
35 nuage2 = generate_gauss(3, 4, 2)
36 nuage = [(a, b, 0) for a, b in nuage1] + [(a, b, 1) for a, b in nuage2]
38 import matplotlib.pyplot as plt
39 from ensae_teaching_cs.helpers.matplotlib_helper_xyz import scatter_xy_id
40 fig, ax = scatter_xy_id(nuage, title="example with random observations",
41 legend={0: "c0", 1: "c1"})
42 plt.show()
44 The error ``ValueError: Unknown projection '3d'`` is raised when the line
45 ``from mpl_toolkits.mplot3d import Axes3D`` is missing.
46 """
47 curves = {}
48 for x, y, zid in xy_id:
49 if zid not in curves:
50 curves[zid] = []
51 curves[zid].append((x, y))
53 marker = options.get('marker', 'o')
54 if ax is None:
55 import matplotlib.pyplot as plt
56 fig, ax = plt.subplots(figsize=options.get('figsize', None))
57 else:
58 fig = None
59 color = 0
60 legs = []
61 for k in sorted(curves):
62 v = curves[k]
63 f, = ax.plot([_[0] for _ in v], [_[1] for _ in v],
64 marker, color=colors_definition[color][1])
65 if legend is not None:
66 legs.append((f, legend.get(k, "unknown:" + str(k))))
67 color += 1
69 if len(legs) > 0:
70 ax.legend([_[0] for _ in legs], [_[1] for _ in legs])
72 if "xlabel" in options:
73 ax.set_xlabel(options["xlabel"])
74 if "ylabel" in options:
75 ax.set_ylabel(options["ylabel"])
76 if "title" in options:
77 ax.set_title(options["title"])
78 return fig, ax
81def scatter_xyc(points, smooth=0, div=10, ax=None, **options):
82 """
83 Draws a 2D graph (X,Y, color), the color is chosen based on a value *f(x,y)*
84 The function requires :epkg:`matplotlib` and :epkg:`scipy`.
86 @param points (x,y, z=f(x,y) )
87 @param smooth applies n times a smoothing I * M (convolutional)
88 @param div number of divisions for axis
89 @param options others options: xlabel, ylabel, title, figsize (if ax is None)
90 @param ax :epkg:`matplotlib` axis
91 @return fig, ax (fig is None if ax was sent to the function)
93 .. plot::
94 :include-source:
96 import random
97 def generate_gauss(x, y, sigma, N=1000):
98 res = []
99 for i in range(N):
100 u = random.gauss(0, 1)
101 a = sigma * u + x
102 b = sigma * random.gauss(0, 1) + y + u
103 res.append((a, b))
104 return res
105 def f(a, b):
106 return (a ** 2 + b ** 2) ** 0.5
107 nuage1 = generate_gauss(0, 0, 3)
108 nuage2 = generate_gauss(3, 4, 2)
109 nuage = [(a, b, f(a, b)) for a, b in nuage1] + [(a, b, f(a, b)) for a, b in nuage2]
110 import matplotlib.pyplot as plt
111 from ensae_teaching_cs.helpers.matplotlib_helper_xyz import scatter_xyc
112 fig, ax = scatter_xyc(nuage, title="example with random observations")
113 plt.show()
115 The error ``ValueError: Unknown projection '3d'`` is raised when the line
116 ``from mpl_toolkits.mplot3d import Axes3D`` is missing.
117 """
118 if ax is None:
119 import matplotlib.pyplot as plt
120 fig, ax = plt.subplots(figsize=options.get('figsize', None))
121 else:
122 fig = None
124 x = [_[0] for _ in points]
125 y = [_[1] for _ in points]
126 z = [_[2] for _ in points]
128 tri = Triangulation(x, y)
130 plt.tricontour(tri, z, 15, linewidths=0.5, colors='k')
131 plt.tricontourf(tri, z, 15, cmap=plt.cm.rainbow,
132 norm=Normalize(vmax=numpy.abs(z).max(), vmin=-numpy.abs(z).max()))
133 plt.colorbar(ax=ax)
134 ax.scatter(x, y, c='b', s=5, zorder=10)
135 ax.set_xlim(min(x), max(x))
136 ax.set_ylim(min(y), max(y))
138 if "xlabel" in options:
139 ax.set_xlabel(options["xlabel"])
140 if "ylabel" in options:
141 ax.set_ylabel(options["ylabel"])
142 if "title" in options:
143 ax.set_title(options["title"])
144 return fig, ax
147def scatter_xyz(points, smooth=0, div=100, ax=None, **options):
148 """
149 Draws a 3D graph (X, Y, Z).
150 The function requires :epkg:`matplotlib`
151 and :epkg:`scipy`.
153 @param points (x,y, z=f(x,y) )
154 @param div number of divisions for axis
155 @param smooth applies n times a smoothing I * M (convolutional)
156 @param ax existing graph to plot on (can be None)
157 @param options others options: xlabel, ylabel, zlabel, title, elev, angle, figsize (if ax is None)
158 - elev, angle: see `view_init <http://matplotlib.org/mpl_toolkits/mplot3d/api.html>`_)
159 @return fig, ax (fig is None if ax was sent to the function)
161 If *ax is None*, axes are created like::
163 fig = plt.figure(figsize=options.get('figsize', None))
164 ax = fig.gca(projection='3d')
166 .. plot::
167 :include-source:
169 import random
170 def generate_gauss(x, y, sigma, N=1000):
171 res = []
172 for i in range(N):
173 u = random.gauss(0, 1)
174 a = sigma * u + x
175 b = sigma * random.gauss(0, 1) + y + u
176 res.append((a, b))
177 return res
178 def f(a, b):
179 return (a ** 2 + b ** 2) ** 0.5
180 nuage1 = generate_gauss(0, 0, 3)
181 nuage2 = generate_gauss(3, 4, 2)
182 nuage = [(a, b, f(a, b)) for a, b in nuage1] + [(a, b, f(a, b)) for a, b in nuage2]
184 import matplotlib.pyplot as plt
185 from ensae_teaching_cs.helpers.matplotlib_helper_xyz import scatter_xyz
186 fig, ax = scatter_xyz(nuage, title="example with random observations")
187 plt.show()
189 The error ``ValueError: Unknown projection '3d'`` is raised when the line
190 ``from mpl_toolkits.mplot3d import Axes3D`` is missing.
191 """
192 x = [_[0] for _ in points]
193 y = [_[1] for _ in points]
194 z = [_[2] for _ in points]
196 if ax is None:
197 import matplotlib.pyplot as plt
198 from mpl_toolkits.mplot3d import Axes3D
199 if Axes3D is None:
200 raise ImportError("Unable to import mpl_toolkits.mplot3d")
201 fig = plt.figure(figsize=options.get('figsize', None))
202 ax = fig.add_subplot(projection='3d')
203 else:
204 fig = None
206 elev = options.get("elev", 50)
207 angle = options.get("angle", 45)
208 ax.view_init(elev, angle)
210 tri = Triangulation(x, y)
211 ax.plot_trisurf(x, y, z, triangles=tri.triangles, cmap='autumn')
213 if "xlabel" in options:
214 ax.set_xlabel(options["xlabel"])
215 if "ylabel" in options:
216 ax.set_ylabel(options["ylabel"])
217 if "zlabel" in options:
218 ax.set_ylabel(options["zlabel"])
219 if "title" in options:
220 ax.set_title(options["title"])
221 return fig, ax