test_epoch_dataarray_accessor.py

import tempfile
from importlib.metadata import version

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pytest
import xarray as xr
from matplotlib.animation import PillowWriter
from matplotlib.container import BarContainer
from packaging.version import Version

import sdf_xarray as sdfxr
import sdf_xarray.plotting as sxp
from sdf_xarray import SDFPreprocess, download

mpl.use("Agg")

# TODO Remove this once the new kwarg options are fully implemented
if Version(version("xarray")) >= Version("2025.8.0"):
    xr.set_options(use_new_combine_kwarg_defaults=True)

TEST_FILES_DIR_1D = download.fetch_dataset("test_files_1D")
TEST_FILES_DIR_2D_MW = download.fetch_dataset("test_files_2D_moving_window")
TEST_FILES_DIR_3D = download.fetch_dataset("test_files_3D")


@pytest.fixture
def subplots():
    fig, ax = plt.subplots()
    yield (fig, ax)
    plt.close(fig)


def test_animation_accessor():
    array = xr.DataArray(
        [1, 2, 3],
        dims=["x"],
        coords={"x": [0, 1, 2]},
        attrs={"long_name": "Test Array", "units": "m"},
    )
    assert hasattr(array, "epoch")
    assert hasattr(array.epoch, "animate")


@pytest.mark.parametrize(
    ("xrlib", "params"),
    [
        (
            xr,
            {
                "compat": "no_conflicts",
                "join": "outer",
                "preprocess": SDFPreprocess(),
            },
        ),
        (sdfxr, {}),
    ],
)
def test_animate_headless(xrlib, params):
    with xrlib.open_mfdataset(TEST_FILES_DIR_1D.glob("*.sdf"), **params) as ds:
        anim = ds["Derived_Number_Density_electron"].epoch.animate()

        # Specify a custom writable temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_file_path = f"{temp_dir}/output.gif"
            try:
                anim.save(temp_file_path, writer=PillowWriter(fps=2))
            except Exception as e:
                pytest.fail(f"animate().save() failed in headless mode: {e}")


def test_get_frame_title_no_optional_params():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = "Time = 5.47e-14 [s]"
        result = sxp.get_frame_title(data, 0)
        assert expected_result == result


def test_get_frame_title_sdf_name():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = "Time = 5.47e-14 [s], 0000.sdf"
        result = sxp.get_frame_title(data, 0, display_sdf_name=True)
        assert expected_result == result


def test_get_frame_title_custom_title():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = "Test Title, Time = 5.47e-14 [s]"
        result = sxp.get_frame_title(data, 0, title_custom="Test Title")
        assert expected_result == result


def test_get_frame_title_custom_title_and_sdf_name():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = "Test Title, Time = 5.47e-14 [s], 0000.sdf"
        result = sxp.get_frame_title(
            data, 0, display_sdf_name=True, title_custom="Test Title"
        )
        assert expected_result == result


def test_get_frame_title_Z_Grid_mid():
    with xr.open_dataset(TEST_FILES_DIR_3D / "0001.sdf") as ds:
        data = ds["Derived_Number_Density_Electron"]
        expected_result = "Z = 3.91e-07 [m]"
        result = sxp.get_frame_title(data, 0, t="Z_Grid_mid")
        assert expected_result == result


def test_calculate_window_boundaries_1D():
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        data = ds["Derived_Number_Density_electron"][:, :, 50]
        expected_result = np.array(
            [[0, 1], [0.49, 1.49], [0.99, 1.99], [1.49, 2.49], [1.99, 2.99]]
        )
        result = sxp.calculate_window_boundaries(data)
        assert result == pytest.approx(expected_result, abs=0.1)


def test_calculate_window_boundaries_2D():
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = np.array(
            [[0, 1], [0.49, 1.49], [0.99, 1.99], [1.49, 2.49], [1.99, 2.99]]
        )
        result = sxp.calculate_window_boundaries(data)
        assert result == pytest.approx(expected_result, abs=0.1)


def test_calculate_window_boundaries_1D_xlim():
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        data = ds["Derived_Number_Density_electron"][:, :, 50]
        expected_result = np.array(
            [[0.1, 0.9], [0.59, 1.39], [1.09, 1.89], [1.59, 2.39], [2.09, 2.89]]
        )
        result = sxp.calculate_window_boundaries(data, xlim=(0.1, 0.9))
        assert result == pytest.approx(expected_result, abs=0.1)


def test_calculate_window_boundaries_2D_xlim():
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        data = ds["Derived_Number_Density_electron"]
        expected_result = np.array(
            [[0.1, 0.9], [0.59, 1.39], [1.09, 1.89], [1.59, 2.39], [2.09, 2.89]]
        )
        result = sxp.calculate_window_boundaries(data, xlim=(0.1, 0.9))
        assert result == pytest.approx(expected_result, abs=0.1)


def test_compute_global_limits():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        result_min, result_max = sxp.compute_global_limits(
            ds["Derived_Number_Density_electron"]
        )
        expected_result_min = 8.07e19
        expected_result_max = 1.17e20
        assert result_min == pytest.approx(expected_result_min, abs=1e18)
        assert result_max == pytest.approx(expected_result_max, abs=1e19)


def test_compute_global_limits_percentile():
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        result_min, result_max = sxp.compute_global_limits(
            ds["Derived_Number_Density_electron"], 40, 45
        )
        expected_result_min = 9.84e19
        expected_result_max = 9.94e19
        assert result_min == pytest.approx(expected_result_min, abs=1e18)
        assert result_max == pytest.approx(expected_result_max, abs=1e18)


def test_compute_global_limits_NaNs():
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        result_min, result_max = sxp.compute_global_limits(
            ds["Derived_Number_Density_electron"]
        )
        expected_result_min = 5.51e-1
        expected_result_max = 2.70
        assert result_min == pytest.approx(expected_result_min, abs=1e-2)
        assert result_max == pytest.approx(expected_result_max, abs=1e-1)


def test_epoch_plot_simple_1d_dataset(subplots):
    with xr.open_mfdataset(
        TEST_FILES_DIR_1D.glob("*.sdf"),
        compat="no_conflicts",
        join="outer",
        preprocess=SDFPreprocess(),
    ) as ds:
        _, ax = subplots
        ds["Derived_Number_Density_electron"].isel(time=0).epoch.plot(ax=ax)

        assert len(ax.lines) == 1
        assert ax.get_xlabel() == "X [m]"


def test_epoch_plot_simple_2d_dataset(subplots):
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        _, ax = subplots
        ds["Derived_Number_Density_electron"].isel(time=0).epoch.plot(ax=ax)

        assert len(ax.collections) > 0
        assert ax.get_xlabel() == "X [m]"
        assert ax.get_ylabel() == "Y [m]"


def test_epoch_plot_simple_3d_dataset_slice(subplots):
    with xr.open_dataset(TEST_FILES_DIR_3D / "0001.sdf") as ds:
        _, ax = subplots
        ds["Derived_Number_Density_Electron"].isel(Z_Grid_mid=0).epoch.plot(ax=ax)

        assert len(ax.collections) > 0
        assert ax.get_xlabel() == "X [m]"
        assert ax.get_ylabel() == "Y [m]"


def test_epoch_plot_flips_axis_order_for_2d_data(subplots):
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        _, ax = subplots
        ds["Derived_Number_Density_electron"].isel(time=0).epoch.plot(ax=ax)

        assert ax.get_xlabel() == "X [m]"
        assert ax.get_ylabel() == "Y [m]"


def test_epoch_plot_flips_axis_order_for_2d_data_with_additional_params(subplots):
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        _, ax = subplots
        ds["Derived_Number_Density_electron"].isel(time=0).epoch.plot(
            ax=ax,
            xlim=(0.5, 1.0),
            ylim=(0.0, 0.5),
        )

        assert ax.get_xlabel() == "X [m]"
        assert ax.get_ylabel() == "Y [m]"
        assert ax.get_xlim() == pytest.approx((0.5, 1.0), abs=1e-2)
        assert ax.get_ylim() == pytest.approx((0.0, 0.5), abs=1e-2)


def test_epoch_plot_flips_axis_order_for_2d_data_but_not_when_time_dim_present(
    subplots,
):
    with xr.open_mfdataset(
        TEST_FILES_DIR_2D_MW.glob("*.sdf"),
        preprocess=SDFPreprocess(),
        combine="nested",
        compat="no_conflicts",
        join="outer",
    ) as ds:
        _, ax = subplots
        plot = ds["Derived_Number_Density_electron"].epoch.plot(ax=ax)

        assert type(plot[2]) is BarContainer