-
Notifications
You must be signed in to change notification settings - Fork 5
/
test_plot.py
52 lines (35 loc) · 1.49 KB
/
test_plot.py
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
import os
from tempfile import TemporaryDirectory
import shutil
import pytest
import numpy as np
import lbmpy.plot as plt
from lbmpy.scenarios import create_lid_driven_cavity
@pytest.mark.skipif(shutil.which('ffmpeg') is None, reason="ffmpeg not available")
def test_animation():
ldc = create_lid_driven_cavity((10, 10), relaxation_rate=1.8)
def run_vec():
ldc.run(100)
return ldc.velocity[:, :, :]
def run_scalar():
ldc.run(100)
return ldc.density[:, :]
plt.clf()
plt.cla()
with TemporaryDirectory() as tmp_dir:
ani = plt.vector_field_magnitude_animation(run_vec, interval=1, frames=2)
ani.save(os.path.join(tmp_dir, "animation1.avi"))
ani = plt.vector_field_animation(run_vec, interval=1, frames=2, rescale=True)
ani.save(os.path.join(tmp_dir, "animation2.avi"))
ani = plt.vector_field_animation(run_vec, interval=1, frames=2, rescale=False)
ani.save(os.path.join(tmp_dir, "animation3.avi"))
ani = plt.scalar_field_animation(run_scalar, interval=1, frames=2, rescale=True)
ani.save(os.path.join(tmp_dir, "animation4.avi"))
ani = plt.scalar_field_animation(run_scalar, interval=1, frames=2, rescale=False)
ani.save(os.path.join(tmp_dir, "animation5.avi"))
ani = plt.surface_plot_animation(run_scalar, frames=2)
ani.save(os.path.join(tmp_dir, "animation6.avi"))
def test_plot():
arr = np.ones([3, 3, 2])
plt.multiple_scalar_fields(arr)
plt.show()