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 

9 

10 

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`. 

15 

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) 

21 

22 .. plot:: 

23 :include-source: 

24 

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] 

37 

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() 

43 

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)) 

52 

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 

68 

69 if len(legs) > 0: 

70 ax.legend([_[0] for _ in legs], [_[1] for _ in legs]) 

71 

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 

79 

80 

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`. 

85 

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) 

92 

93 .. plot:: 

94 :include-source: 

95 

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() 

114 

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 

123 

124 x = [_[0] for _ in points] 

125 y = [_[1] for _ in points] 

126 z = [_[2] for _ in points] 

127 

128 tri = Triangulation(x, y) 

129 

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)) 

137 

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 

145 

146 

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`. 

152 

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) 

160 

161 If *ax is None*, axes are created like:: 

162 

163 fig = plt.figure(figsize=options.get('figsize', None)) 

164 ax = fig.gca(projection='3d') 

165 

166 .. plot:: 

167 :include-source: 

168 

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] 

183 

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() 

188 

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] 

195 

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.gca(projection='3d') 

203 else: 

204 fig = None 

205 

206 elev = options.get("elev", 50) 

207 angle = options.get("angle", 45) 

208 ax.view_init(elev, angle) 

209 

210 tri = Triangulation(x, y) 

211 ax.plot_trisurf(x, y, z, triangles=tri.triangles, cmap='autumn') 

212 

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