Skip to content

Commit

Permalink
migrate test_vr_orientation.py to pytest-mpl
Browse files Browse the repository at this point in the history
  • Loading branch information
neutrinoceros committed Sep 12, 2022
1 parent 4d3a3f8 commit 1a43611
Showing 1 changed file with 86 additions and 78 deletions.
164 changes: 86 additions & 78 deletions yt/visualization/volume_rendering/tests/test_vr_orientation.py
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

0 comments on commit 1a43611

Please sign in to comment.