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()