diff --git a/examples/example_data/diffusion_data_example.h5 b/examples/example_data/diffusion_data_example.h5 new file mode 100644 index 0000000..78baba3 Binary files /dev/null and b/examples/example_data/diffusion_data_example.h5 differ diff --git a/examples/example_data/vanadium_data_example.h5 b/examples/example_data/vanadium_data_example.h5 new file mode 100644 index 0000000..0c7534a Binary files /dev/null and b/examples/example_data/vanadium_data_example.h5 differ diff --git a/examples/experiment_example.ipynb b/examples/experiment_example.ipynb new file mode 100644 index 0000000..565e16c --- /dev/null +++ b/examples/experiment_example.ipynb @@ -0,0 +1,48 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "51b7b8be", + "metadata": {}, + "outputs": [], + "source": [ + "from easydynamics.experiment import Experiment\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eb91e49a", + "metadata": {}, + "outputs": [], + "source": [ + "vanadium_experiment=Experiment(\"Vanadium\")\n", + "vanadium_experiment.load_hdf5(filename=\"example_data/vanadium_data_example.h5\")\n", + "\n", + "vanadium_experiment.plot_data()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "newdynamics", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/easydynamics/experiment/__init__.py b/src/easydynamics/experiment/__init__.py new file mode 100644 index 0000000..154fc7e --- /dev/null +++ b/src/easydynamics/experiment/__init__.py @@ -0,0 +1,7 @@ +from .experiment import ( + Experiment, +) + +__all__ = [ + "Experiment", +] diff --git a/src/easydynamics/experiment/experiment.py b/src/easydynamics/experiment/experiment.py new file mode 100644 index 0000000..e38740f --- /dev/null +++ b/src/easydynamics/experiment/experiment.py @@ -0,0 +1,137 @@ +from typing import Optional, Union + +import plopp as pp +import scipp as sc +from easyscience.job.experiment import ExperimentBase +from scipp.io import load_hdf5 as sc_load_hdf5 +from scipp.io import save_hdf5 as sc_save_hdf5 + + +class Experiment(ExperimentBase): + """ + Holds data from an experiment as a sc.DataArray along with metadata. + """ + + def __init__( + self, + name: str, + data: Optional[Union[sc.DataArray, str]] = None, + *args, + **kwargs, + ): + super().__init__(name, *args, **kwargs) + + if data is None: + self._data: Optional[sc.DataArray] = None + elif isinstance(data, str): + self.load_hdf5(filename=data) + elif isinstance(data, sc.DataArray): + self._data = data + else: + raise TypeError( + f"Data must be a sc.DataArray or a filename string, not {type(data).__name__}" + ) + + def load_hdf5(self, filename: str, name: Optional[str] = None): + """ + Load data from an HDF5 file. + + Args: + filename (str): Path to the HDF5 file. + """ + if not isinstance(filename, str): + raise TypeError(f"Filename must be a string, not {type(filename).__name__}") + + if name is not None: + if not isinstance(name, str): + raise TypeError(f"Name must be a string, not {type(name).__name__}") + self.name = name + + # TODO: Add checks of dimensions etc. I'm not yet sure what dimensions I want to allow, so for now I trust that the data is valid. + loaded_data = sc_load_hdf5(filename) + if not isinstance(loaded_data, sc.DataArray): + raise TypeError( + f"Loaded data must be a sc.DataArray, not {type(loaded_data).__name__}" + ) + self._data = loaded_data + + def save_hdf5(self, filename: Optional[str] = None): + """Save the dataset to HDF5. + + Args: + filename (str): Path to the output HDF5 file. + """ + + if filename is None: + filename = f"{self.name}.h5" + + if not isinstance(filename, str): + raise TypeError(f"Filename must be a string, not {type(filename).__name__}") + + if self._data is None: + raise ValueError("No data to save.") + + import os + + dir_name = os.path.dirname(filename) + if dir_name: + os.makedirs(dir_name, exist_ok=True) + + sc_save_hdf5(self._data, filename) + + def remove_data(self): + """Remove the dataset from the experiment.""" + self._data = None + + @property + def data(self) -> Optional[sc.DataArray]: + """Get the dataset associated with this experiment.""" + return self._data + + @data.setter + def data(self, value: sc.DataArray): + """Set the dataset associated with this experiment.""" + if not isinstance(value, sc.DataArray): + raise TypeError(f"Data must be a sc.DataArray, not {type(value).__name__}") + self._data = value + + def plot_data(self): + """Plot the dataset using plopp.""" + + if self._data is None: + raise ValueError("No data to plot. Please load data first.") + + if not self._in_notebook(): + raise RuntimeError( + "plot_data() can only be used in a Jupyter notebook environment." + ) + + from IPython.display import display + + fig = pp.plot(self._data.transpose(), title=f"{self.name}") + display(fig) + + @staticmethod + def _in_notebook(): + try: + from IPython import get_ipython + + shell = get_ipython().__class__.__name__ + if shell == "ZMQInteractiveShell": + return True # Jupyter notebook or JupyterLab + elif shell == "TerminalInteractiveShell": + return False # Terminal IPython + else: + return False + except (NameError, ImportError): + return False # Standard Python (no IPython) + + def __repr__(self) -> str: + return f"Experiment `{self.name}` with data: {self._data}" + + def __copy__(self) -> "Experiment": + """Return a copy of the object.""" + temp = self.as_dict(skip=["unique_name"]) + new_obj = self.__class__.from_dict(temp) + new_obj.data = self.data.copy() if self.data is not None else None + return new_obj diff --git a/tests/unit_tests/experiment/test_experiment.py b/tests/unit_tests/experiment/test_experiment.py new file mode 100644 index 0000000..f17ff4b --- /dev/null +++ b/tests/unit_tests/experiment/test_experiment.py @@ -0,0 +1,277 @@ +from copy import copy +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import scipp as sc + +from easydynamics.experiment import Experiment + + +class TestExperiment: + @pytest.fixture + def experiment(self): + Q = sc.linspace("Q", 0.5, 1.5, num=10, unit="1/Angstrom") + energy = sc.linspace("energy", -5, 5, num=11, unit="meV") + values = sc.array(dims=["Q", "energy"], values=np.ones((10, 11))) + data = sc.DataArray(data=values, coords={"Q": Q, "energy": energy}) + + experiment = Experiment(name="test_experiment", data=data) + return experiment + + def test_init_array(self, experiment): + "Test initialization with a Scipp DataArray" + # WHEN THEN EXPECT + assert experiment.name == "test_experiment" + assert isinstance(experiment._data, sc.DataArray) + assert "Q" in experiment._data.dims + assert "energy" in experiment._data.dims + assert experiment._data.sizes["Q"] == 10 + assert experiment._data.sizes["energy"] == 11 + assert sc.identical( + experiment._data.data, + sc.array(dims=["Q", "energy"], values=np.ones((10, 11))), + ) + + def test_init_string(self, tmp_path): + "Test initialization with a filename string - should load the file" + # WHEN + Q = sc.linspace("Q", 0.5, 1.5, num=10, unit="1/Angstrom") + energy = sc.linspace("energy", -5, 5, num=11, unit="meV") + values = sc.array(dims=["Q", "energy"], values=np.ones((10, 11))) + data = sc.DataArray(data=values, coords={"Q": Q, "energy": energy}) + + filename = tmp_path / "test_experiment.h5" + sc.io.save_hdf5(data, filename) + + # THEN + experiment = Experiment(name="loaded_experiment", data=str(filename)) + + # EXPECT + assert experiment.name == "loaded_experiment" + assert isinstance(experiment._data, sc.DataArray) + assert "Q" in experiment._data.dims + assert "energy" in experiment._data.dims + assert experiment._data.sizes["Q"] == 10 + assert experiment._data.sizes["energy"] == 11 + assert sc.identical( + experiment._data.data, + sc.array(dims=["Q", "energy"], values=np.ones((10, 11))), + ) + + def test_init_no_data(self): + "Test initialization with no data" + # WHEN + experiment = Experiment(name="empty_experiment") + + # THEN EXPECT + assert experiment.name == "empty_experiment" + assert experiment._data is None + + def test_init_invalid_data(self): + "Test initialization with invalid data type" + # WHEN / THEN EXPECT + with pytest.raises(TypeError): + Experiment(name="invalid_experiment", data=123) + + def test_load_hdf5(self, tmp_path, experiment): + "Test loading data from an HDF5 file. First use scipp to save data to a file, then load it using the method." + # WHEN + # First create a file to load from + filename = tmp_path / "test.h5" + data_to_save = experiment.data + sc.io.save_hdf5(data_to_save, filename) + + # THEN + new_experiment = Experiment("new_experiment") + new_experiment.load_hdf5(str(filename), name="loaded_data") + loaded_data = new_experiment.data + + # EXPECT + assert sc.identical(data_to_save, loaded_data) + assert new_experiment.name == "loaded_data" + + def test_load_hdf5_invalid_name_raises(self, experiment): + "Test loading data from an HDF5 file, giving the Experiment an invalid name" + # WHEN / THEN EXPECT + with pytest.raises(TypeError): + experiment.load_hdf5("some_file.h5", name=123) + + def test_load_hdf5_invalid_filename_raises(self, experiment): + "Test loading data from an HDF5 file with an invalid filename" + # WHEN / THEN EXPECT + with pytest.raises(TypeError): + experiment.load_hdf5(123) + + def test_load_hdf5_invalid_file_raises(self, experiment): + "Test loading data from a non-existent HDF5 file" + # WHEN / THEN EXPECT + with pytest.raises(OSError): + experiment.load_hdf5("non_existent_file.h5") + + def test_save_hdf5(self, tmp_path, experiment): + "Test saving data to an HDF5 file. Load the saved file using scipp and compare to the original data." + # WHEN THEN + filename = tmp_path / "saved_data.h5" + experiment.save_hdf5(str(filename)) + + # EXPECT + loaded_data = sc.io.load_hdf5(str(filename)) + original_data = experiment.data + assert sc.identical(original_data, loaded_data) + + def test_save_hdf5_default_filename(self, tmp_path, experiment, monkeypatch): + "Test saving data to an HDF5 file with default filename" + # WHEN + monkeypatch.chdir(tmp_path) + + # THEN + experiment.save_hdf5() + + # EXPECT + expected_filename = tmp_path / f"{experiment.name}.h5" + loaded_data = sc.io.load_hdf5(str(expected_filename)) + original_data = experiment.data + assert sc.identical(original_data, loaded_data) + + def test_save_hdf5_no_data_raises(self): + "Test saving data to an HDF5 file when no data is present in the experiment" + # WHEN + experiment = Experiment(name="no_data_experiment") + + # THEN EXPECT + with pytest.raises(ValueError): + experiment.save_hdf5("should_fail.h5") + + def test_save_hdf5_invalid_filename_raises(self, experiment): + "Test saving data to an HDF5 file with an invalid filename" + # WHEN / THEN EXPECT + with pytest.raises(TypeError): + experiment.save_hdf5(123) + + def test_remove_data(self, experiment): + "Test removing data from the experiment" + # WHEN + experiment.remove_data() + + # THEN EXPECT + assert experiment._data is None + + def test_data_setter_raises_type_error(self, experiment): + "Test setting data to an invalid type raises TypeError" + # WHEN THEN EXPECT + with pytest.raises(TypeError): + experiment.data = 123 + + def test_repr(self, experiment): + # WHEN + repr_str = repr(experiment) + + # THEN EXPECT + assert ( + repr_str == f"Experiment `{experiment.name}` with data: {experiment._data}" + ) + + def test_copy_experiment(self, experiment): + "Test copying an Experiment object. The copied object should have the same attributes but be a different object in memory." + # WHEN + copied_experiment = copy(experiment) + + # THEN EXPECT + assert copied_experiment.name == experiment.name + assert sc.identical(copied_experiment.data, experiment.data) + assert copied_experiment is not experiment + assert copied_experiment.data is not experiment.data + + def test_plot_data_success(self, experiment): + "Test plotting data successfully when in notebook environment" + # WHEN + with ( + patch.object(Experiment, "_in_notebook", return_value=True), + patch("plopp.plot") as mock_plot, + patch("IPython.display.display") as mock_display, + ): + mock_fig = MagicMock() + mock_plot.return_value = mock_fig + + # THEN + experiment.plot_data() + + # EXPECT + mock_plot.assert_called_once() + args, kwargs = mock_plot.call_args + assert sc.identical(args[0], experiment._data.transpose()) + assert kwargs["title"] == f"{experiment.name}" + mock_display.assert_called_once_with(mock_fig) + + def test_plot_data_no_data_raises(self): + "Test plotting data raises ValueError when no data is present" + # WHEN + experiment = Experiment(name="empty_experiment") + + # THEN EXPECT + with pytest.raises(ValueError, match="No data to plot"): + experiment.plot_data() + + def test_plot_data_not_in_notebook_raises(self, experiment): + "Test plotting data raises RuntimeError when not in notebook environment" + # WHEN + with patch.object(Experiment, "_in_notebook", return_value=False): + # THEN EXPECT + with pytest.raises( + RuntimeError, + match="plot_data\\(\\) can only be used in a Jupyter notebook environment", + ): + experiment.plot_data() + + def test_in_notebook_returns_true_for_jupyter(self, monkeypatch): + """Should return True when IPython shell is ZMQInteractiveShell (Jupyter).""" + + # WHEN + class ZMQInteractiveShell: + __name__ = "ZMQInteractiveShell" + + # THEN + monkeypatch.setattr("IPython.get_ipython", lambda: ZMQInteractiveShell()) + + # EXPECT + assert Experiment._in_notebook() is True + + def test_in_notebook_returns_false_for_terminal_ipython(self, monkeypatch): + """Should return False when IPython shell is TerminalInteractiveShell.""" + + # WHEN + class TerminalInteractiveShell: + __name__ = "TerminalInteractiveShell" + + # THEN + + monkeypatch.setattr("IPython.get_ipython", lambda: TerminalInteractiveShell()) + + # EXPECT + assert Experiment._in_notebook() is False + + def test_in_notebook_returns_false_for_unknown_shell(self, monkeypatch): + """Should return False when IPython shell type is unrecognized.""" + + # WHEN + class UnknownShell: + __name__ = "UnknownShell" + + # THEN + monkeypatch.setattr("IPython.get_ipython", lambda: UnknownShell()) + # EXPECT + assert Experiment._in_notebook() is False + + def test_in_notebook_returns_false_when_no_ipython(self, monkeypatch): + """Should return False when IPython is not installed or available.""" + + # WHEN + def raise_import_error(*args, **kwargs): + raise ImportError + + # THEN + monkeypatch.setattr("builtins.__import__", raise_import_error) + + # EXPECT + assert Experiment._in_notebook() is False