wolfhece 2.2.15__py3-none-any.whl → 2.2.17__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,807 @@
1
+ from matplotlib import pyplot as plt
2
+ from matplotlib.axes import Axes
3
+ from matplotlib.figure import Figure
4
+ import numpy as np
5
+
6
+ from ..wolf_array import header_wolf
7
+ from ..wolfresults_2D import q_splitting, u_splitting, splitting_rule
8
+
9
+ class Mesh2D(header_wolf):
10
+
11
+ def __init__(self, src_header:header_wolf):
12
+
13
+ self.set_origin(src_header.origx, src_header.origy)
14
+ self.set_resolution(src_header.dx, src_header.dy)
15
+ self.set_translation(src_header.translx, src_header.transly)
16
+ self.shape = src_header.shape
17
+ self._factor = None
18
+
19
+ def plot_cells(self, ax:Axes=None, transpose:bool= False, color='black', **kwargs):
20
+ """ Plot the grid of the mesh.
21
+ """
22
+
23
+ if ax is None:
24
+ fig, ax = plt.subplots()
25
+ else:
26
+ fig = ax.figure
27
+
28
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
29
+
30
+ if transpose:
31
+
32
+ # plot the grid of the mesh in a transposed way
33
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1):
34
+ ax.plot([ymin, ymax], [x, x], color=color, **kwargs)
35
+
36
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby + 1):
37
+ ax.plot([y, y], [xmin, xmax], color=color, **kwargs)
38
+
39
+ self.set_aspect_labels_matrice(ax=ax, **kwargs)
40
+
41
+ else:
42
+
43
+ # plot the grid of the mesh
44
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby + 1):
45
+ ax.plot([xmin, xmax], [y, y], color=color, **kwargs)
46
+
47
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1):
48
+ ax.plot([x, x], [ymin, ymax], color=color, **kwargs)
49
+
50
+ self.set_aspect_labels(ax=ax, **kwargs)
51
+
52
+ return fig, ax
53
+
54
+ def plot_center_cells(self, ax:Axes=None, color='black', linestyle='--', **kwargs):
55
+ """ Plot lines centered to the cells.
56
+ """
57
+
58
+ if ax is None:
59
+ fig, ax = plt.subplots()
60
+ else:
61
+ fig = ax.figure
62
+
63
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
64
+
65
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
66
+ ax.plot([xmin, xmax], [y, y], color=color, linestyle=linestyle, **kwargs)
67
+
68
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
69
+ ax.plot([x, x], [ymin, ymax], color=color, linestyle=linestyle, **kwargs)
70
+
71
+ self.set_aspect_labels(ax=ax, **kwargs)
72
+ return fig, ax
73
+
74
+ def set_ticks_as_dxdy(self, ax:Axes=None, **kwargs):
75
+ """ Set the ticks of the axis as the dx and dy of the mesh.
76
+ """
77
+
78
+ if ax is None:
79
+ fig, ax = plt.subplots()
80
+ else:
81
+ fig = ax.figure
82
+
83
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
84
+
85
+ x_ticks = ['ox + tx', r'+$\Delta x$'] + [f'+{i}$\Delta x$' for i in range(2,self.nbx+1)]
86
+ y_ticks = ['oy + ty', r'+$\Delta y$'] + [f'{+i}$\Delta y$' for i in range(2,self.nby+1)]
87
+
88
+ ax.set_xticks(np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1))
89
+ ax.set_yticks(np.linspace(ymin, ymax, endpoint=True, num=self.nby + 1))
90
+ ax.set_xticklabels(x_ticks)
91
+ ax.set_yticklabels(y_ticks)
92
+
93
+ return fig, ax
94
+
95
+ def set_ticks_as_matrice(self, ax:Axes=None, Fortran_type:bool = True, **kwargs):
96
+ """ Set the ticks of the axis as the row and column of a matrice """
97
+
98
+ if ax is None:
99
+ fig, ax = plt.subplots()
100
+ else:
101
+ fig = ax.figure
102
+
103
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
104
+
105
+ if Fortran_type:
106
+ x_ticks = [f'{i}' for i in range(1,self.nbx+1)]
107
+ y_ticks = [f'{i}' for i in range(1,self.nby+1)]
108
+ else:
109
+ x_ticks = [f'{i}' for i in range(self.nbx)]
110
+ y_ticks = [f'{i}' for i in range(self.nby)]
111
+
112
+ ax.set_yticks(np.linspace(xmin+self.dx/2., xmax-self.dx/2., endpoint=True, num=self.nbx))
113
+ ax.set_xticks(np.linspace(ymin+self.dy/2., ymax-self.dy/2., endpoint=True, num=self.nby))
114
+
115
+ x_ticks.reverse()
116
+ ax.set_yticklabels(x_ticks)
117
+ ax.set_xticklabels(y_ticks)
118
+
119
+ self.set_aspect_labels_matrice(ax=ax, **kwargs)
120
+
121
+ return fig, ax
122
+
123
+ def plot_circle_at_centers(self, ax:Axes=None, color='black', radius:float=None, **kwargs):
124
+ """ Plot circles at the center of the cells.
125
+ """
126
+
127
+ if radius is None:
128
+ radius = 0.1 * self.dx
129
+
130
+ if ax is None:
131
+ fig, ax = plt.subplots()
132
+ else:
133
+ fig = ax.figure
134
+
135
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
136
+
137
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
138
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
139
+ circle = plt.Circle((x, y), radius=radius, color=color, **kwargs)
140
+ ax.add_artist(circle)
141
+
142
+ self.set_aspect_labels(ax=ax, **kwargs)
143
+
144
+ return fig, ax
145
+
146
+ def plot_indices_at_centers(self, ax:Axes=None, Fortran_type:bool = True, **kwargs):
147
+ """ Plot the indices of the cells at the center of the cells.
148
+ """
149
+
150
+ if ax is None:
151
+ fig, ax = plt.subplots()
152
+ else:
153
+ fig = ax.figure
154
+
155
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
156
+
157
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
158
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
159
+ i,j = self.xy2ij(x, y)
160
+
161
+ if Fortran_type:
162
+ i+=1
163
+ j+=1
164
+
165
+ ax.text(x, y, f'({i},{j})', horizontalalignment='center', verticalalignment='center', **kwargs)
166
+
167
+ self.set_aspect_labels(ax=ax, **kwargs)
168
+
169
+ return fig, ax
170
+
171
+ def plot_memoryposition_at_centers(self, ax:Axes=None,
172
+ transpose=False,
173
+ Fortran_type:bool = True,
174
+ f_contiguous:bool = True,
175
+ **kwargs):
176
+ """ Plot the position of the cells at the center of the cells.
177
+ """
178
+
179
+ if ax is None:
180
+ fig, ax = plt.subplots()
181
+ else:
182
+ fig = ax.figure
183
+
184
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
185
+
186
+ if transpose:
187
+ k = 0
188
+ if Fortran_type:
189
+ k+=1
190
+
191
+ all_y = list(np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx))
192
+ all_x = list(np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby))
193
+
194
+ all_y.reverse()
195
+
196
+ if f_contiguous:
197
+
198
+ for x in all_x:
199
+ for y in all_y:
200
+
201
+ ax.text(x, y, f'{k}', horizontalalignment='center', verticalalignment='center', **kwargs)
202
+ k+=1
203
+ else:
204
+
205
+ for y in all_y:
206
+ for x in all_x:
207
+
208
+ ax.text(x, y, f'{k}', horizontalalignment='center', verticalalignment='center', **kwargs)
209
+ k+=1
210
+
211
+ self.set_aspect_labels_matrice(ax=ax, **kwargs)
212
+
213
+ else:
214
+
215
+ k = 0
216
+ if Fortran_type:
217
+ k+=1
218
+
219
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
220
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
221
+
222
+ ax.text(x, y, f'{k}', horizontalalignment='center', verticalalignment='center', **kwargs)
223
+ k+=1
224
+
225
+ self.set_aspect_labels(ax=ax, **kwargs)
226
+
227
+ return fig, ax
228
+
229
+ def plot_indices_at_bordersX(self, ax:Axes=None, Fortran_type:bool = True, **kwargs):
230
+ """ Plot the indices of the cells at the borders of the cells.
231
+ """
232
+
233
+ if ax is None:
234
+ fig, ax = plt.subplots()
235
+ else:
236
+ fig = ax.figure
237
+
238
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
239
+
240
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
241
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1):
242
+ i,j = self.xy2ij(x, y)
243
+
244
+ if Fortran_type:
245
+ i+=1
246
+ j+=1
247
+
248
+ ax.text(x, y, f'({i},{j})', horizontalalignment='center', verticalalignment='center', **kwargs)
249
+
250
+ def plot_indices_at_bordersY(self, ax:Axes=None, Fortran_type:bool = True, **kwargs):
251
+ """ Plot the indices of the cells at the borders of the cells.
252
+ """
253
+
254
+ if ax is None:
255
+ fig, ax = plt.subplots()
256
+ else:
257
+ fig = ax.figure
258
+
259
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
260
+
261
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby + 1):
262
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
263
+ i,j = self.xy2ij(x, y)
264
+
265
+ if Fortran_type:
266
+ i+=1
267
+ j+=1
268
+
269
+ ax.text(x, y, f'({i},{j})', horizontalalignment='center', verticalalignment='center', **kwargs)
270
+
271
+ self.set_aspect_labels(ax=ax, **kwargs)
272
+
273
+ return fig, ax
274
+
275
+ def plot_Xarrows_at_center(self, ax:Axes=None,
276
+ randomize:bool=False,
277
+ amplitude:np.ndarray=None,
278
+ color='black', **kwargs):
279
+ """ Plot arrows at the center of the cells.
280
+ """
281
+
282
+ if ax is None:
283
+ fig, ax = plt.subplots()
284
+ else:
285
+ fig = ax.figure
286
+
287
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
288
+
289
+ if amplitude is None:
290
+ amplitude = np.ones((self.nbx, self.nby)) * 0.2 * self.dx
291
+ else:
292
+ amplitude = self.scale_amplitude(amplitude)
293
+
294
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
295
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
296
+
297
+ i,j = self.xy2ij(x, y)
298
+ if randomize:
299
+ dx = np.random.uniform(-0.5, 0.5) * self.dx
300
+ amplitude[i,j] = dx
301
+ else:
302
+ dx = amplitude[i,j]
303
+
304
+ ax.arrow(x - dx/2., y, dx, 0,
305
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
306
+ fc=color, ec=color)
307
+
308
+ self.set_aspect_labels(ax=ax, **kwargs)
309
+
310
+ return fig, ax, amplitude
311
+
312
+ def plot_Yarrows_at_center(self, ax:Axes=None,
313
+ randomize:bool=False,
314
+ amplitude:np.ndarray=None,
315
+ color='black', **kwargs):
316
+ """ Plot arrows at the center of the cells.
317
+ """
318
+
319
+ if ax is None:
320
+ fig, ax = plt.subplots()
321
+ else:
322
+ fig = ax.figure
323
+
324
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
325
+
326
+ if amplitude is None:
327
+ amplitude = np.ones((self.nbx, self.nby)) * 0.2 * self.dy
328
+ else:
329
+ amplitude = self.scale_amplitude(amplitude)
330
+
331
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
332
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
333
+
334
+ i,j = self.xy2ij(x, y)
335
+ if randomize:
336
+ dy = np.random.uniform(-0.5, 0.5) * self.dy
337
+ amplitude[i,j] = dy
338
+ else:
339
+ dy = amplitude[i,j]
340
+
341
+ ax.arrow(x, y - dy/2., 0, dy,
342
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
343
+ fc=color, ec=color)
344
+
345
+ self.set_aspect_labels(ax=ax, **kwargs)
346
+
347
+ return fig, ax, amplitude
348
+
349
+ def plot_Xarrows_at_borders(self, ax:Axes=None,
350
+ randomize:bool=False,
351
+ amplitudeX:np.ndarray=None,
352
+ amplitudeY:np.ndarray=None,
353
+ color='black', **kwargs):
354
+ """ Plot arrows at the borders of the cells.
355
+ """
356
+
357
+ if ax is None:
358
+ fig, ax = plt.subplots()
359
+ else:
360
+ fig = ax.figure
361
+
362
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
363
+
364
+ if amplitudeX is None:
365
+ amplitudeX = np.ones((self.nbx+1, self.nby)) * 0.2 * self.dx
366
+ else:
367
+ if self._factor is not None:
368
+ amplitudeX *= self._factor
369
+
370
+ if amplitudeY is None:
371
+ amplitudeY = np.ones((self.nbx, self.nby+1)) * 0.2 * self.dy
372
+ else:
373
+ if self._factor is not None:
374
+ amplitudeY *= self._factor
375
+
376
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
377
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1):
378
+
379
+ i,j = self.xy2ij(x, y)
380
+ if randomize:
381
+ dx = np.random.uniform(-0.5, 0.5) * self.dx
382
+ amplitudeX[i,j] = dx
383
+ else:
384
+ dx = amplitudeX[i,j]
385
+
386
+ if dx != 0.:
387
+ ax.arrow(x - dx/2., y, dx, 0,
388
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
389
+ fc=color, ec=color)
390
+
391
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby+1):
392
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
393
+
394
+ i,j = self.xy2ij(x, y)
395
+ if randomize:
396
+ dx = np.random.uniform(-0.5, 0.5) * self.dy
397
+ amplitudeY[i,j] = dx
398
+ else:
399
+ dx = amplitudeY[i,j]
400
+
401
+ if dx != 0.:
402
+ ax.arrow(x - dx, y, dx, 0,
403
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
404
+ fc=color, ec=color)
405
+
406
+ self.set_aspect_labels(ax=ax, **kwargs)
407
+
408
+ return fig, ax, amplitudeX, amplitudeY
409
+
410
+ def plot_Yarrows_at_borders(self, ax:Axes=None,
411
+ randomize:bool=False,
412
+ amplitudeX:np.ndarray=None,
413
+ amplitudeY:np.ndarray=None,
414
+ color='black', **kwargs):
415
+ """ Plot arrows at the borders of the cells.
416
+ """
417
+
418
+ if ax is None:
419
+ fig, ax = plt.subplots()
420
+ else:
421
+ fig = ax.figure
422
+
423
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
424
+
425
+ if amplitudeX is None:
426
+ amplitudeX = np.ones((self.nbx+1, self.nby)) * 0.2 * self.dx
427
+ else:
428
+ if self._factor is not None:
429
+ amplitudeX *= self._factor
430
+
431
+ if amplitudeY is None:
432
+ amplitudeY = np.ones((self.nbx, self.nby+1)) * 0.2 * self.dy
433
+ else:
434
+ if self._factor is not None:
435
+ amplitudeY *= self._factor
436
+
437
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby + 1):
438
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
439
+
440
+ i,j = self.xy2ij(x, y)
441
+ if randomize:
442
+ dy = np.random.uniform(-0.5, 0.5) * self.dy
443
+ amplitudeY[i,j] = dy
444
+ else:
445
+ dy = amplitudeY[i,j]
446
+
447
+ if dy != 0.:
448
+ ax.arrow(x, y - dy/2., 0, dy,
449
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
450
+ fc=color, ec=color)
451
+
452
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
453
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx + 1):
454
+
455
+ i,j = self.xy2ij(x, y)
456
+ if randomize:
457
+ dy = np.random.uniform(-0.5, 0.5) * self.dx
458
+ amplitudeX[i,j] = dy
459
+ else:
460
+ dy = amplitudeX[i,j]
461
+
462
+ if dy != 0.:
463
+ ax.arrow(x, y - dy/2., 0., dy,
464
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
465
+ fc=color, ec=color)
466
+
467
+ self.set_aspect_labels(ax=ax, **kwargs)
468
+
469
+ return fig, ax, amplitudeX, amplitudeY
470
+
471
+ def plot_normal_arrows_at_borders(self, ax:Axes=None, color='black', **kwargs):
472
+ """ Plot arrows at the borders of the cells.
473
+ """
474
+
475
+ if ax is None:
476
+ fig, ax = plt.subplots()
477
+ else:
478
+ fig = ax.figure
479
+
480
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
481
+
482
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
483
+ for x in np.linspace(xmin, xmax, endpoint=True, num=self.nbx+1):
484
+ ax.arrow(x - 0.1 * self.dx, y, 0.2 * self.dx, 0,
485
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
486
+ fc=color, ec=color)
487
+
488
+ for y in np.linspace(ymin, ymax, endpoint=True, num=self.nby+1):
489
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
490
+ ax.arrow(x, y - 0.1 * self.dy, 0., 0.2 * self.dy,
491
+ head_width=0.1 * self.dx, head_length=0.1 * self.dy,
492
+ fc=color, ec=color)
493
+
494
+ self.set_aspect_labels(ax=ax, **kwargs)
495
+
496
+ return fig, ax
497
+
498
+ def scale_axes(self, ax:Axes=None, factor:float = 0.1, **kwargs):
499
+ """ Scale the axes of the plot to fit the data.
500
+ """
501
+ # augmenter légèrement la taille visible sans
502
+ # ajouter de ticks
503
+ ticksX = ax.get_xticks()
504
+ ticksY = ax.get_yticks()
505
+ ticklabelsX = ax.get_xticklabels()
506
+ ticklabelsY = ax.get_yticklabels()
507
+
508
+ deltaX = ticksX[-1] - ticksX[0]
509
+ deltaY = ticksY[-1] - ticksY[0]
510
+ ax.set_xlim(ticksX[0] - factor * deltaX, ticksX[-1] + factor * deltaX)
511
+ ax.set_ylim(ticksY[0] - factor * deltaY, ticksY[-1] + factor * deltaY)
512
+
513
+ ax.set_xticks(ticksX)
514
+ ax.set_yticks(ticksY)
515
+ ax.set_xticklabels(ticklabelsX) #, rotation=45)
516
+ ax.set_yticklabels(ticklabelsY) #, rotation=45)
517
+
518
+ # remove up border and right border
519
+ ax.spines['top'].set_visible(False)
520
+ ax.spines['right'].set_visible(False)
521
+
522
+ return ax.figure, ax
523
+
524
+ def plot_reconstructed_values_at_borders(self, ax:Axes=None,
525
+ colors=['green', 'blue', 'red', 'brown'],
526
+ radius:float = None, **kwargs):
527
+ """ Plot 4 small circles on each side of the border.
528
+ """
529
+
530
+ if radius is None:
531
+ radius = 0.02 * self.dx
532
+
533
+ if ax is None:
534
+ fig, ax = plt.subplots()
535
+ else:
536
+ fig = ax.figure
537
+
538
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
539
+
540
+ deltay = self.dy / 6.
541
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
542
+ x = xmin
543
+ for i in range(4):
544
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i*deltay),
545
+ radius=radius, color=colors[-i-1], **kwargs))
546
+ x = xmax
547
+ for i in range(4):
548
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i*deltay),
549
+ radius=radius, color=colors[-i-1], **kwargs))
550
+
551
+ for x in np.linspace(xmin + self.dx, xmax - self.dx, endpoint=True, num=self.nbx-1):
552
+ for i in range(4):
553
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i*deltay),
554
+ radius=radius, color=colors[-i-1], **kwargs))
555
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i*deltay),
556
+ radius=radius, color=colors[-i-1], **kwargs))
557
+
558
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
559
+ y = ymin
560
+ for i in range(4):
561
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i*deltay, y+.1*self.dy),
562
+ radius=radius, color=colors[i], **kwargs))
563
+ y = ymax
564
+ for i in range(4):
565
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i*deltay, y-.1*self.dy),
566
+ radius=radius, color=colors[i], **kwargs))
567
+
568
+ for y in np.linspace(ymin + self.dy, ymax - self.dy, endpoint=True, num=self.nby-1):
569
+ for i in range(4):
570
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i*deltay, y+.1*self.dy),
571
+ radius=radius, color=colors[i], **kwargs))
572
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i*deltay, y-.1*self.dy),
573
+ radius=radius, color=colors[i], **kwargs))
574
+
575
+ self.set_aspect_labels(ax=ax, **kwargs)
576
+
577
+ def plot_splitted_values_at_borders(self, ax:Axes=None,
578
+ qx:np.ndarray=None,
579
+ qy:np.ndarray=None,
580
+ colors=['green', 'blue', 'red', 'brown'],
581
+ radius:float = None, **kwargs):
582
+ """ Plot 4 small circles on each side of the border.
583
+ """
584
+
585
+ split_x = self.zeros_bordersX()
586
+ split_y = self.zeros_bordersY()
587
+
588
+ for i in range(1, self.nbx):
589
+ for j in range(self.nby):
590
+ split_x[i,j] = splitting_rule(qx[i-1,j], qx[i,j])
591
+ for i in range(self.nbx):
592
+ for j in range(1, self.nby):
593
+ split_y[i,j] = splitting_rule(qy[i,j-1], qy[i,j])
594
+
595
+ if radius is None:
596
+ radius = 0.02 * self.dx
597
+
598
+ if ax is None:
599
+ fig, ax = plt.subplots()
600
+ else:
601
+ fig = ax.figure
602
+
603
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
604
+
605
+ deltay = self.dy / 6.
606
+
607
+ # X borders
608
+ for y in np.linspace(ymin + self.dy/2., ymax - self.dy/2., endpoint=True, num=self.nby):
609
+ x = xmin
610
+ i,j = self.xy2ij(x, y)
611
+ for i_unk in range(4):
612
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i_unk*deltay),
613
+ radius=radius, color=colors[-i_unk-1], **kwargs))
614
+ x = xmax
615
+ i,j = self.xy2ij(x, y)
616
+ for i_unk in range(4):
617
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i_unk*deltay),
618
+ radius=radius, color=colors[-i_unk-1], **kwargs))
619
+
620
+ for x in np.linspace(xmin + self.dx, xmax - self.dx, endpoint=True, num=self.nbx-1):
621
+ i,j = self.xy2ij(x, y)
622
+
623
+ pond = split_x[i,j]
624
+ if pond == 0.5:
625
+ for i_unk in range(4):
626
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i_unk*deltay),
627
+ radius=radius, color=colors[-i_unk-1], **kwargs))
628
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i_unk*deltay),
629
+ radius=radius, color=colors[-i_unk-1], **kwargs))
630
+ elif pond == 1.0:
631
+ # left is upstream
632
+ for i_unk in [1,2]:
633
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i_unk*deltay),
634
+ radius=radius, color=colors[-i_unk-1], **kwargs))
635
+ for i_unk in [0,3]:
636
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i_unk*deltay),
637
+ radius=radius, color=colors[-i_unk-1], **kwargs))
638
+ elif pond == 0.0:
639
+ # right is upstream
640
+ for i_unk in [0,3]:
641
+ ax.add_artist(plt.Circle((x-.1*self.dx, y - 1.5*deltay + i_unk*deltay),
642
+ radius=radius, color=colors[-i_unk-1], **kwargs))
643
+ for i_unk in [1,2]:
644
+ ax.add_artist(plt.Circle((x+.1*self.dx, y - 1.5*deltay + i_unk*deltay),
645
+ radius=radius, color=colors[-i_unk-1], **kwargs))
646
+
647
+ # Y borders
648
+ for x in np.linspace(xmin + self.dx/2., xmax - self.dx/2., endpoint=True, num=self.nbx):
649
+ y = ymin
650
+ i,j = self.xy2ij(x, y)
651
+ for i_unk in range(4):
652
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y+.1*self.dy),
653
+ radius=radius, color=colors[i_unk], **kwargs))
654
+ y = ymax
655
+ i,j = self.xy2ij(x, y)
656
+ for i_unk in range(4):
657
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y-.1*self.dy),
658
+ radius=radius, color=colors[i_unk], **kwargs))
659
+
660
+ for y in np.linspace(ymin + self.dy, ymax - self.dy, endpoint=True, num=self.nby-1):
661
+ i,j = self.xy2ij(x, y)
662
+ pond = split_y[i,j]
663
+ if pond == 0.5:
664
+ for i_unk in range(4):
665
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y+.1*self.dy),
666
+ radius=radius, color=colors[i_unk], **kwargs))
667
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y-.1*self.dy),
668
+ radius=radius, color=colors[i_unk], **kwargs))
669
+ elif pond == 1.0:
670
+ for i_unk in [1,2]:
671
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y-.1*self.dy),
672
+ radius=radius, color=colors[i_unk], **kwargs))
673
+ for i_unk in [0,3]:
674
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y+.1*self.dy),
675
+ radius=radius, color=colors[i_unk], **kwargs))
676
+ elif pond == 0.0:
677
+ for i_unk in [0,3]:
678
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y-.1*self.dy),
679
+ radius=radius, color=colors[i_unk], **kwargs))
680
+ for i_unk in [1,2]:
681
+ ax.add_artist(plt.Circle((x - 1.5*deltay + i_unk*deltay, y+.1*self.dy),
682
+ radius=radius, color=colors[i_unk], **kwargs))
683
+
684
+
685
+ self.set_aspect_labels(ax=ax, **kwargs)
686
+
687
+ def set_aspect_labels(self, ax:Axes=None, **kwargs):
688
+ """ Set the aspect of the plot to be equal.
689
+ """
690
+ if ax is None:
691
+ fig, ax = plt.subplots()
692
+ else:
693
+ fig = ax.figure
694
+
695
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
696
+
697
+ ax.set_aspect('equal')
698
+ ax.set_xlim(xmin, xmax)
699
+ ax.set_ylim(ymin, ymax)
700
+ ax.set_xlabel('X (m)')
701
+ ax.set_ylabel('Y (m)')
702
+ return fig, ax
703
+
704
+ def set_aspect_labels_matrice(self, ax:Axes=None, **kwargs):
705
+ """ Set the aspect of the plot to be equal.
706
+ """
707
+ if ax is None:
708
+ fig, ax = plt.subplots()
709
+ else:
710
+ fig = ax.figure
711
+
712
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
713
+
714
+ ax.set_aspect('equal')
715
+ ax.set_ylim(xmin, xmax)
716
+ ax.set_xlim(ymin, ymax)
717
+ ax.set_xlabel('columns')
718
+ ax.set_ylabel('rows')
719
+
720
+ #set x ais on the upper side
721
+ ax.xaxis.set_ticks_position('top')
722
+ ax.xaxis.set_label_position('top')
723
+ return fig, ax
724
+
725
+ def zeros(self):
726
+ """ Return a 2D array of zeros with the shape of the mesh.
727
+ """
728
+ return np.zeros((self.nbx, self.nby))
729
+
730
+ def ones(self):
731
+ """ Return a 2D array of ones with the shape of the mesh.
732
+ """
733
+ return np.ones((self.nbx, self.nby))
734
+
735
+ def zeros_bordersX(self):
736
+ """ Return a 2D array of zeros with the shape of the mesh + 1 in x direction.
737
+ """
738
+ return np.zeros((self.nbx+1, self.nby))
739
+
740
+ def zeros_bordersY(self):
741
+ """ Return a 2D array of zeros with the shape of the mesh + 1 in y direction.
742
+ """
743
+ return np.zeros((self.nbx, self.nby+1))
744
+
745
+ def ones_bordersX(self):
746
+ """ Return a 2D array of ones with the shape of the mesh + 1 in x direction.
747
+ """
748
+ return np.ones((self.nbx+1, self.nby))
749
+
750
+ def ones_bordersY(self):
751
+ """ Return a 2D array of ones with the shape of the mesh + 1 in y direction.
752
+ """
753
+ return np.ones((self.nbx, self.nby+1))
754
+
755
+ def apply_splitting_X(self, q:np.ndarray):
756
+ """ Apply the splitting rule to the X direction.
757
+ """
758
+ q_borderX = self.zeros_bordersX()
759
+ for i in range(1,self.nbx):
760
+ for j in range(self.nby):
761
+ q_borderX[i,j] = q_splitting(q[i-1,j], q[i,j])
762
+ return q_borderX
763
+
764
+ def apply_splitting_Y(self, q:np.ndarray):
765
+ """ Apply the splitting rule to the Y direction.
766
+ """
767
+ q_borderY = self.zeros_bordersY()
768
+ for i in range(self.nbx):
769
+ for j in range(1,self.nby):
770
+ q_borderY[i,j] = q_splitting(q[i,j-1], q[i,j])
771
+ return q_borderY
772
+
773
+ def scale_amplitude(self, amplitude:np.ndarray, factor:float = None):
774
+ """ Scale the amplitude of the arrows.
775
+ """
776
+
777
+ if self._factor is not None:
778
+ return amplitude * self._factor
779
+
780
+ if factor is not None:
781
+ return amplitude * factor
782
+
783
+ factor = min(self.dx, self.dy) * .5 / np.max(amplitude)
784
+ self._factor = factor
785
+
786
+ return amplitude * factor
787
+
788
+ def plot_outside_domain(self, ax:Axes=None, color='black', **kwargs):
789
+ """ Plot a hashed zone outside of the domain.
790
+ """
791
+ if ax is None:
792
+ fig, ax = plt.subplots()
793
+ else:
794
+ fig = ax.figure
795
+
796
+ [xmin, xmax], [ymin, ymax] = self.get_bounds()
797
+
798
+ # Create a rectangle outside the domain
799
+ rect = plt.Rectangle((xmin-self.dx, ymin-self.dy), xmax - xmin + 2.*self.dx, ymax - ymin + 2.*self.dy,
800
+ color=color, alpha=0.2, hatch='//', **kwargs)
801
+ ax.add_patch(rect)
802
+
803
+ rect2 = plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
804
+ color='white', alpha=1., hatch='', **kwargs)
805
+ ax.add_patch(rect2)
806
+
807
+ return fig, ax