-
Notifications
You must be signed in to change notification settings - Fork 280
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
migrate test_vr_orientation.py to pytest-mpl
- Loading branch information
1 parent
4d3a3f8
commit 1a43611
Showing
1 changed file
with
86 additions
and
78 deletions.
There are no files selected for viewing
164 changes: 86 additions & 78 deletions
164
yt/visualization/volume_rendering/tests/test_vr_orientation.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,101 +1,109 @@ | ||
import os | ||
import tempfile | ||
|
||
import matplotlib as mpl | ||
import matplotlib.pyplot as plt | ||
import numpy as np | ||
from nose.plugins.attrib import attr | ||
import pytest | ||
|
||
from yt.testing import ANSWER_TEST_TAG, fake_vr_orientation_test_ds | ||
from yt.utilities.answer_testing.framework import ( | ||
GenericImageTest, | ||
VRImageComparisonTest, | ||
) | ||
from yt.testing import fake_vr_orientation_test_ds | ||
from yt.visualization.volume_rendering.api import ( | ||
ColorTransferFunction, | ||
Scene, | ||
create_volume_source, | ||
off_axis_projection, | ||
) | ||
|
||
|
||
@attr(ANSWER_TEST_TAG) | ||
def test_orientation(): | ||
ds = fake_vr_orientation_test_ds() | ||
def scene_to_mpl_figure(scene): | ||
"""helper function to convert a scene image rendering to matplotlib | ||
so we can rely on pytest-mpl to compare images | ||
""" | ||
tmpfd, tmpname = tempfile.mkstemp(suffix=".png") | ||
os.close(tmpfd) | ||
scene.save(tmpname, sigma_clip=1.0) | ||
image = mpl.image.imread(tmpname) | ||
os.remove(tmpname) | ||
|
||
sc = Scene() | ||
fig, ax = plt.subplots() | ||
ax.set(aspect="equal") | ||
ax.imshow(image) | ||
return fig | ||
# zlib.compress(image.dumps()) | ||
|
||
vol = create_volume_source(ds, field=("gas", "density")) | ||
sc.add_source(vol) | ||
|
||
tf = vol.transfer_function | ||
tf = ColorTransferFunction((0.1, 1.0)) | ||
tf.sample_colormap(1.0, 0.01, colormap="coolwarm") | ||
tf.sample_colormap(0.8, 0.01, colormap="coolwarm") | ||
tf.sample_colormap(0.6, 0.01, colormap="coolwarm") | ||
tf.sample_colormap(0.3, 0.01, colormap="coolwarm") | ||
class TestOrientation: | ||
@classmethod | ||
def setup_class(cls): | ||
cls.ds = fake_vr_orientation_test_ds() | ||
|
||
n_frames = 1 | ||
orientations = [[-0.3, -0.1, 0.8]] | ||
cls.scene = Scene() | ||
|
||
theta = np.pi / n_frames | ||
test_name = "vr_orientation" | ||
vol = create_volume_source(cls.ds, field=("gas", "density")) | ||
cls.scene.add_source(vol) | ||
|
||
for lens_type, decimals in [("perspective", 12), ("plane-parallel", 2)]: | ||
# set a much lower precision for plane-parallel tests, see | ||
# https://github.com/yt-project/yt/issue/3069 | ||
# https://github.com/yt-project/yt/pull/3068 | ||
# https://github.com/yt-project/yt/pull/3294 | ||
frame = 0 | ||
cls._last_lense_type = None | ||
|
||
@classmethod | ||
def set_camera(cls, lens_type): | ||
# this method isn't thread-safe | ||
# if lens_type == cls._last_lense_type: | ||
# return | ||
|
||
cam = sc.add_camera(ds, lens_type=lens_type) | ||
cls._last_lense_type = lens_type | ||
|
||
cam = cls.scene.add_camera(cls.ds, lens_type=lens_type) | ||
cam.resolution = (1000, 1000) | ||
cam.position = ds.arr(np.array([-4.0, 0.0, 0.0]), "code_length") | ||
cam.position = cls.ds.arr(np.array([-4.0, 0.0, 0.0]), "code_length") | ||
cam.switch_orientation( | ||
normal_vector=[1.0, 0.0, 0.0], north_vector=[0.0, 0.0, 1.0] | ||
) | ||
cam.set_width(ds.domain_width * 2.0) | ||
desc = "%s_%04d" % (lens_type, frame) | ||
test1 = VRImageComparisonTest(sc, ds, desc, decimals) | ||
test1.answer_name = test_name | ||
yield test1 | ||
|
||
for _ in range(n_frames): | ||
frame += 1 | ||
center = ds.arr([0, 0, 0], "code_length") | ||
cam.yaw(theta, rot_center=center) | ||
desc = "yaw_%s_%04d" % (lens_type, frame) | ||
test2 = VRImageComparisonTest(sc, ds, desc, decimals) | ||
test2.answer_name = test_name | ||
yield test2 | ||
|
||
for _ in range(n_frames): | ||
frame += 1 | ||
theta = np.pi / n_frames | ||
center = ds.arr([0, 0, 0], "code_length") | ||
cam.pitch(theta, rot_center=center) | ||
desc = "pitch_%s_%04d" % (lens_type, frame) | ||
test3 = VRImageComparisonTest(sc, ds, desc, decimals) | ||
test3.answer_name = test_name | ||
yield test3 | ||
|
||
for _ in range(n_frames): | ||
frame += 1 | ||
theta = np.pi / n_frames | ||
center = ds.arr([0, 0, 0], "code_length") | ||
cam.roll(theta, rot_center=center) | ||
desc = "roll_%s_%04d" % (lens_type, frame) | ||
test4 = VRImageComparisonTest(sc, ds, desc, decimals) | ||
test4.answer_name = test_name | ||
yield test4 | ||
|
||
center = [0.5, 0.5, 0.5] | ||
width = [1.0, 1.0, 1.0] | ||
|
||
for i, orientation in enumerate(orientations): | ||
cam.set_width(cls.ds.domain_width * 2.0) | ||
cls.camera = cam | ||
|
||
@pytest.mark.parametrize("lens_type", ["perspective", "plane-parallel"]) | ||
@pytest.mark.mpl_image_compare | ||
def test_lense_type(self, lens_type): | ||
# TODO(4118): | ||
# set a much lower precision for plane-parallel tests, see | ||
# https://github.com/yt-project/yt/issue/3069 | ||
# https://github.com/yt-project/yt/pull/3068 | ||
# https://github.com/yt-project/yt/pull/3294 | ||
self.set_camera(lens_type) | ||
return scene_to_mpl_figure(self.scene) | ||
|
||
@pytest.mark.mpl_image_compare | ||
def test_yaw(self): | ||
self.set_camera("plane-parallel") | ||
center = self.ds.arr([0, 0, 0], "code_length") | ||
self.camera.yaw(np.pi, rot_center=center) | ||
return scene_to_mpl_figure(self.scene) | ||
|
||
@pytest.mark.mpl_image_compare | ||
def test_pitch(self): | ||
self.set_camera("plane-parallel") | ||
center = self.ds.arr([0, 0, 0], "code_length") | ||
self.camera.pitch(np.pi, rot_center=center) | ||
return scene_to_mpl_figure(self.scene) | ||
|
||
@pytest.mark.mpl_image_compare | ||
def test_roll(self): | ||
self.set_camera("plane-parallel") | ||
center = self.ds.arr([0, 0, 0], "code_length") | ||
self.camera.roll(np.pi, rot_center=center) | ||
return scene_to_mpl_figure(self.scene) | ||
|
||
@pytest.mark.mpl_image_compare | ||
def test_off_axis_projection(self): | ||
image = off_axis_projection( | ||
ds, center, orientation, width, 512, ("gas", "density"), no_ghost=False | ||
self.ds, | ||
center=[0.5, 0.5, 0.5], | ||
normal_vector=[-0.3, -0.1, 0.8], | ||
width=[1.0, 1.0, 1.0], | ||
resolution=512, | ||
item=("gas", "density"), | ||
no_ghost=False, | ||
) | ||
|
||
def offaxis_image_func(filename_prefix): | ||
return image.write_image(filename_prefix) # noqa: B023 | ||
|
||
test5 = GenericImageTest(ds, offaxis_image_func, decimals) | ||
test5.prefix = f"oap_orientation_{i}" | ||
test5.answer_name = test_name | ||
yield test5 | ||
fig, ax = plt.subplots() | ||
ax.imshow(image) | ||
return fig |