From 742e2da90a02418d51900b9bc08ab2b3273390fc Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Fri, 22 Mar 2024 21:35:01 +0900 Subject: [PATCH] =?UTF-8?q?=E2=9C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/arpes/io.py | 15 +++++++++++++-- src/arpes/widgets.py | 4 ++-- src/arpes/xarray_extensions.py | 11 +++++------ tests/test_xarray_extensions.py | 18 +++++++++++++++++- 4 files changed, 37 insertions(+), 11 deletions(-) diff --git a/src/arpes/io.py b/src/arpes/io.py index 6dccc2ea..ce4e9c32 100644 --- a/src/arpes/io.py +++ b/src/arpes/io.py @@ -104,12 +104,14 @@ def load_data( return load_scan(desc, **kwargs) -DATA_EXAMPLES = { +DATA_EXAMPLES: dict[str, tuple[str, str]] = { "cut": ("ALG-MC", "cut.fits"), "map": ("example_data", "fermi_surface.nc"), "photon_energy": ("example_data", "photon_energy.nc"), "nano_xps": ("example_data", "nano_xps.nc"), "temperature_dependence": ("example_data", "temperature_dependence.nc"), + "cut2": ("SPD", "example_itx_data.itx"), + "map2": ("IF_UMCS", "BLGr_GK_example_xy_data.xy"), } @@ -117,7 +119,7 @@ def load_example_data(example_name: str = "cut") -> xr.Dataset: """Provides sample data for executable documentation. Args: - example_name: [TODO:description] + example_name: (cut, map, photon_energy, nano_xps, temperature_dependence) Returns: [TODO:description] @@ -131,6 +133,7 @@ def load_example_data(example_name: str = "cut") -> xr.Dataset: ) location, example = DATA_EXAMPLES[example_name] + logger.debug(f"location:{location}") file = Path(__file__).parent / "example_data" / example return load_data(file=file, location=location) @@ -157,6 +160,14 @@ def nano_xps(self) -> xr.Dataset: def temperature_dependence(self) -> xr.Dataset: return load_example_data("temperature_dependence") + @property + def cut2(self) -> xr.Dataset: + return load_example_data("cut2") + + @property + def map2(self) -> xr.Dataset: + return load_example_data("map2") + example_data = ExampleData() diff --git a/src/arpes/widgets.py b/src/arpes/widgets.py index 116fcd0b..c8c7e3a6 100644 --- a/src/arpes/widgets.py +++ b/src/arpes/widgets.py @@ -35,7 +35,7 @@ from collections.abc import Sequence from functools import wraps from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger -from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, TypeVar +from typing import TYPE_CHECKING, Any, Literal, ParamSpec, TypeAlias, TypeVar import matplotlib as mpl import matplotlib.pyplot as plt @@ -757,7 +757,7 @@ def kspace_tool( overplot_bz: Callable[[Axes], None] | list[Callable[[Axes], None]] | None = None, bounds: dict[MOMENTUM, tuple[float, float]] | None = None, resolution: dict | None = None, - coords: dict[str, NDArray[np.float_] | xr.DataArray] | None = None, + coords: dict[Literal["kp", "kx", "ky", "kz"], NDArray[np.float_]] | None = None, **kwargs: Incomplete, ) -> CurrentContext: """A utility for assigning coordinate offsets using a live momentum conversion. diff --git a/src/arpes/xarray_extensions.py b/src/arpes/xarray_extensions.py index cc2c4474..c2873227 100644 --- a/src/arpes/xarray_extensions.py +++ b/src/arpes/xarray_extensions.py @@ -770,13 +770,14 @@ def dshape(self) -> dict[Hashable, int]: @property def scan_name(self) -> str: - """[TODO:summary]. + """Return scan name. Args: self ([TODO:type]): [TODO:description] - Returns: - [TODO:description] + Returns: (str) + If "scan" or "file" is set in attrs, return the file name. + If they are not set, return "id" if "id" is set. """ for option in ["scan", "file"]: if option in self._obj.attrs: @@ -784,9 +785,7 @@ def scan_name(self) -> str: id_code = self._obj.attrs.get("id") - if id_code is None: - return "No ID" - return str(id_code) + return str(id_code) if id_code is not None else "No ID" @property def label(self) -> str: diff --git a/tests/test_xarray_extensions.py b/tests/test_xarray_extensions.py index 1929713f..05f6a93e 100644 --- a/tests/test_xarray_extensions.py +++ b/tests/test_xarray_extensions.py @@ -30,6 +30,18 @@ def hv_map() -> xr.Dataset: return example_data.photon_energy +@pytest.fixture() +def dataset_cut2() -> xr.Dataset: + """A fixture for loading Dataset.""" + return example_data.cut2 + + +@pytest.fixture() +def dataarray_cut2() -> xr.DataArray: + """A fixture for loading Dataset.""" + return example_data.cut2.spectrum + + class TestforProperties: """Test class for Array Dataset properties.""" @@ -298,6 +310,9 @@ def test_switch_energy_notation( with pytest.raises(RuntimeError) as e: hv_map.S.switch_energy_notation() assert str(e.value) == "Not impremented yet." + with pytest.raises(RuntimeError) as e: + hv_map.spectrum.S.switch_energy_notation() + assert str(e.value) == "Not impremented yet." def test_spectrum_type(self, dataarray_cut: xr.DataArray) -> None: """Test spectrum_type.""" @@ -305,9 +320,10 @@ def test_spectrum_type(self, dataarray_cut: xr.DataArray) -> None: del dataarray_cut.attrs["spectrum_type"] assert dataarray_cut.S.spectrum_type == "cut" - def test_label(self, dataarray_cut: xr.DataArray) -> None: + def test_label(self, dataarray_cut: xr.DataArray, dataarray_cut2: xr.DataArray) -> None: """Test scan_name.""" assert dataarray_cut.S.label == "cut.fits" + assert dataarray_cut2.S.scan_name == "2" class TestGeneralforDataArray: