Source code for xbout.tests.test_animate

import pytest
from matplotlib import pyplot as plt
import numpy as np
import xarray as xr

from xbout import open_boutdataset
from .utils_for_tests import create_bout_ds_list

from animatplot.blocks import Pcolormesh, Line


[docs] @pytest.fixture def create_test_file(tmp_path_factory): # Create temp dir for output of animate1D/2D save_dir = tmp_path_factory.mktemp("test_data") # Generate some test data ds_list, file_list = create_bout_ds_list( "BOUT.dmp", nxpe=3, nype=3, syn_data_type="linear" ) for ds, file_name in zip(ds_list, file_list): ds.to_netcdf(save_dir.joinpath(file_name)) with pytest.warns(UserWarning): ds = open_boutdataset(save_dir.joinpath("BOUT.dmp.*.nc")) # Open test data return save_dir, ds
[docs] class TestAnimate: """ Set of tests to check whether animate1D() and animate2D() are running properly and PillowWriter is saving each animation correctly """
[docs] def test_animate2D(self, create_test_file): save_dir, ds = create_test_file animation = ds["n"].isel(x=1).bout.animate2D(save_as="%s/testyz" % save_dir) assert len(animation.blocks) == 1 block = animation.blocks[0] assert isinstance(block, Pcolormesh) assert block.ax.get_xlabel() == "y" assert block.ax.get_ylabel() == "z" plt.close() animation = ds["n"].isel(y=2).bout.animate2D(save_as="%s/testxz" % save_dir) assert len(animation.blocks) == 1 block = animation.blocks[0] assert isinstance(block, Pcolormesh) assert block.ax.get_xlabel() == "x" assert block.ax.get_ylabel() == "z" plt.close() animation = ds["n"].isel(z=3).bout.animate2D(save_as="%s/testxy" % save_dir) assert len(animation.blocks) == 1 block = animation.blocks[0] assert isinstance(block, Pcolormesh) assert block.ax.get_xlabel() == "x" assert block.ax.get_ylabel() == "y" plt.close()
[docs] @pytest.mark.parametrize( "controls", [ ("both", False), ("timeline", False), ("toggle", False), ("", False), (None, False), ("foo", True), ], ) def test_animate2D_controls_arg(self, create_test_file, controls): controls, expect_error = controls save_dir, ds = create_test_file if expect_error: with pytest.raises(ValueError, match="Unrecognised value for controls"): animation = ds["n"].isel(x=1).bout.animate2D(controls=controls) else: animation = ds["n"].isel(x=1).bout.animate2D(controls=controls) assert len(animation.blocks) == 1 assert isinstance(animation.blocks[0], Pcolormesh) if controls in ["both", "timeline"]: assert hasattr(animation, "slider") else: assert not hasattr(animation, "slider") if controls in ["both", "toggle"]: assert hasattr(animation, "button") else: assert not hasattr(animation, "button") plt.close()
[docs] def test_animate1D(self, create_test_file): save_dir, ds = create_test_file animation = ds["n"].isel(y=2, z=0).bout.animate1D(save_as="%s/test" % save_dir) assert len(animation.blocks) == 1 assert isinstance(animation.blocks[0], Line) plt.close()
[docs] @pytest.mark.parametrize( "controls", [ ("both", False), ("timeline", False), ("toggle", False), ("", False), (None, False), ("foo", True), ], ) def test_animate1D_controls_arg(self, create_test_file, controls): controls, expect_error = controls save_dir, ds = create_test_file if expect_error: with pytest.raises(ValueError, match="Unrecognised value for controls"): animation = ds["n"].isel(y=2, z=0).bout.animate1D(controls=controls) else: animation = ds["n"].isel(y=2, z=0).bout.animate1D(controls=controls) assert len(animation.blocks) == 1 assert isinstance(animation.blocks[0], Line) if controls in ["both", "timeline"]: assert hasattr(animation, "slider") else: assert not hasattr(animation, "slider") if controls in ["both", "toggle"]: assert hasattr(animation, "button") else: assert not hasattr(animation, "button") plt.close()
[docs] def test_animate_list(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)] ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_1d_default(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(y=2, z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)] ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Line) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_1d_multiline(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(y=2, z=3).bout.animate_list( [["n", "T"], ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)] ) assert len(animation.blocks) == 4 assert isinstance(animation.blocks[0], Line) assert isinstance(animation.blocks[1], Line) assert isinstance(animation.blocks[2], Pcolormesh) assert isinstance(animation.blocks[3], Line) # check there were actually 3 subplots assert len([x for x in plt.gcf().get_axes() if x.get_xlabel() != ""]) == 3 plt.close()
[docs] def test_animate_list_animate_over(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(t=2), ds["n"].isel(y=1, z=2)], animate_over="x" ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_save_as(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], save_as="%s/test" % save_dir, ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_fps(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], fps=42 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) assert animation.timeline.fps == 42 plt.close()
[docs] def test_animate_list_nrows(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], nrows=2 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_ncols(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], ncols=3 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_not_enough_nrowsncols(self, create_test_file): save_dir, ds = create_test_file with pytest.raises(ValueError): ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], nrows=2, ncols=1 )
[docs] @pytest.mark.skip(reason="test data for plot_poloidal needs more work") def test_animate_list_poloidal_plot(self, create_test_file): save_dir, ds = create_test_file metadata = ds.metadata metadata["ixseps1"] = 2 metadata["ixseps2"] = 4 metadata["jyseps1_1"] = 1 metadata["jyseps2_1"] = 2 metadata["ny_inner"] = 3 metadata["jyseps1_2"] = 4 metadata["jyseps2_2"] = 5 from ..geometries import apply_geometry from ..utils import _set_attrs_on_all_vars ds = _set_attrs_on_all_vars(ds, "metadata", metadata) nx = ds.metadata["nx"] ny = ds.metadata["ny"] R = xr.DataArray( np.ones([nx, ny]) * np.linspace(0, 1, nx)[:, np.newaxis], dims=["x", "y"] ) Z = xr.DataArray( np.ones([nx, ny]) * np.linspace(0, 1, ny)[np.newaxis, :], dims=["x", "y"] ) ds["psixy"] = R ds["Rxy"] = R ds["Zxy"] = Z ds = apply_geometry(ds, "toroidal") animation = ds.isel(zeta=3).bout.animate_list( ["n", ds["T"].isel(zeta=3), ds["n"].isel(theta=1, zeta=2)], poloidal_plot=True, ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_subplots_adjust(self, create_test_file): save_dir, ds = create_test_file with pytest.warns(UserWarning): animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], subplots_adjust={"hspace": 4, "wspace": 5}, ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_vmin(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], vmin=-0.1 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_vmin_list(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], vmin=[0.0, 0.1, 0.2] ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_vmax(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], vmax=1.1 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_vmax_list(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], vmax=[1.0, 1.1, 1.2] ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_logscale(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], logscale=True ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_logscale_float(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], logscale=1.0e-2 ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_logscale_list(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], logscale=[True, 1.0e-2, False], ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) plt.close()
[docs] def test_animate_list_titles_list(self, create_test_file): save_dir, ds = create_test_file animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], titles=["a", None, "b"] ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert animation.blocks[0].ax.title.get_text() == "a" assert isinstance(animation.blocks[1], Pcolormesh) assert animation.blocks[1].ax.title.get_text() == "T" assert isinstance(animation.blocks[2], Line) assert animation.blocks[2].ax.title.get_text() == "b" plt.close()
[docs] @pytest.mark.parametrize( "controls", [ ("both", False), ("timeline", False), ("toggle", False), ("", False), (None, False), ("foo", True), ], ) def test_animate_list_controls_arg(self, create_test_file, controls): controls, expect_error = controls save_dir, ds = create_test_file if expect_error: with pytest.raises(ValueError, match="Unrecognised value for controls"): animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], controls=controls ) else: animation = ds.isel(z=3).bout.animate_list( ["n", ds["T"].isel(x=2), ds["n"].isel(y=1, z=2)], controls=controls ) assert len(animation.blocks) == 3 assert isinstance(animation.blocks[0], Pcolormesh) assert isinstance(animation.blocks[1], Pcolormesh) assert isinstance(animation.blocks[2], Line) if controls in ["both", "timeline"]: assert hasattr(animation, "slider") else: assert not hasattr(animation, "slider") if controls in ["both", "toggle"]: assert hasattr(animation, "button") else: assert not hasattr(animation, "button") plt.close()