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
|
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
figure, axes = plt.subplots()
class Plotter:
def __init__(self):
self.plt = None
self.colorbar = None
self.patch = None
def clear(self):
if self.patch:
self.patch.remove()
if self.colorbar:
self.colorbar.remove()
if self.plt:
self.plt.remove()
def plot(self, model, normalize=True, density=False, save_path=''):
self.clear()
axes.set_title('Velocity field (normalized)')
plt.suptitle(f'Avg mass source per grid point = {model.avg_error()}')
plt.xlabel('X')
plt.ylabel('Y')
u, v = model.u, model.v
if normalize:
factor = np.sqrt(u ** 2 + v ** 2)
u = u / factor
v = v / factor
x, y = np.meshgrid(
np.linspace(0, model.domain_size[1], model.shape[1]),
np.linspace(0, model.domain_size[0], model.shape[0]),
)
density = density or int((max(model.domain_size) / model.step) / 40)
plt.contourf(x, y, model.p)
# self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
# TODO: allow using streamplot
self.plt = plt.quiver(
x[::density, ::density],
y[::density, ::density],
u[::density, ::density],
v[::density, ::density],
)
self.colorbar = plt.colorbar(label='Pressure')
def save(self, path):
return plt.savefig(path, dpi=300)
def show(self):
return plt.pause(0.0001)
|
