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
|
import numpy as np
import matplotlib.pyplot as plt
from copy import copy
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
shape = (model.p.shape[0] + 1, model.p.shape[1] + 1)
x, y = np.meshgrid(
np.linspace(0, shape[1] * model.step, shape[1]),
np.linspace(0, shape[0] * model.step, shape[0]),
)
u = copy(model.u)
u.resize(shape)
v = copy(model.v)
v.resize(shape)
p = copy(model.p)
p.resize(shape)
print(shape, u.shape, v.shape)
# density = density or int((max(model.domain_size) / model.step) / 40)
plt.contourf(x, y, p)
# self.patch = axes.add_patch(Rectangle((0, 0), *reversed(model.bfs_size), color='gray'))
# TODO: allow using streamplot
self.plt = plt.quiver(
x,
y,
u,
v,
)
self.colorbar = plt.colorbar(label='Pressure')
def save(self, path):
return plt.savefig(path, dpi=300)
def show(self):
return plt.pause(0.0001)
|
