From c3abe3517046ed603a9221de38b22257322d3a51 Mon Sep 17 00:00:00 2001 From: Kimoon Han <98246499+kmnhan@users.noreply.github.com> Date: Thu, 11 Apr 2024 11:58:51 +0900 Subject: [PATCH] fix(interactive): update data load functions used in imagetool --- docs/source/user-guide/io.ipynb | 5 +-- src/erlab/interactive/imagetool/__init__.py | 36 ++++++++++++--------- src/erlab/io/__init__.py | 22 ++++++++----- src/erlab/io/dataloader.py | 10 ------ src/erlab/io/plugins/da30.py | 4 +-- src/erlab/io/utilities.py | 19 +++++++---- 6 files changed, 49 insertions(+), 47 deletions(-) diff --git a/docs/source/user-guide/io.ipynb b/docs/source/user-guide/io.ipynb index c3ce1aab..62814313 100644 --- a/docs/source/user-guide/io.ipynb +++ b/docs/source/user-guide/io.ipynb @@ -654,10 +654,7 @@ "used to access the loader.\n", "\n", "If the ``name`` attribute is prefixed with an underscore, the loader is not registered\n", - "as a loader. This is useful for base classes that are not meant to be used directly. One\n", - "example is :class:`DA30Loader `, which only implements\n", - "the ``load_single`` method and serves as a base class for implementing loaders for\n", - "setups using the Scienta Omicron DA30 hemispherical analyzer with ``SES.exe``.\n", + "as a loader. This is useful for base classes that are not meant to be used directly.\n", "\n", "ARPES data from a single experiment are usually stored in one folder, with files that\n", "look like ``file_0001.h5``, ``file_0002.h5``, etc. If the naming scheme does not deviate\n", diff --git a/src/erlab/interactive/imagetool/__init__.py b/src/erlab/interactive/imagetool/__init__.py index 3e76a92e..3e9f1380 100644 --- a/src/erlab/interactive/imagetool/__init__.py +++ b/src/erlab/interactive/imagetool/__init__.py @@ -36,7 +36,7 @@ from erlab.interactive.utilities import DictMenuBar, copy_to_clipboard if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence import numpy as np import numpy.typing as npt @@ -421,24 +421,24 @@ def _copy_cursor_idx(self): copy_to_clipboard(str(self.slicer_area.array_slicer._indices)) def _open_file(self): - valid_files = { - "xarray HDF5 Files (*.h5)": (xr.load_dataarray, {"engine": "h5netcdf"}), - "ALS BL4.0.3 Raw Data (*.pxt)": (erlab.io.merlin.load, {}), - "ALS BL4.0.3 Live (*.ibw)": (erlab.io.merlin.load_live, {}), - "DA30 Raw Data (*.ibw *.pxt *.zip)": (erlab.io.da30.load, {}), - "SSRL BL5-2 Raw Data (*.h5)": (erlab.io.ssrl52.load, {}), + valid_loaders: dict[str, tuple[Callable, dict]] = { + "xarray HDF5 Files (*.h5)": (erlab.io.load_hdf5, {}), + "ALS BL4.0.3 Raw Data (*.pxt)": (erlab.io.loaders["merlin"].load, {}), + "ALS BL4.0.3 Live (*.ibw)": (erlab.io.loaders["merlin"].load_live, {}), + "DA30 Raw Data (*.ibw *.pxt *.zip)": (erlab.io.loaders["da30"].load, {}), + "SSRL BL5-2 Raw Data (*.h5)": (erlab.io.loaders["ssrl52"].load, {}), "NetCDF Files (*.nc *.nc4 *.cdf)": (xr.load_dataarray, {}), } dialog = QtWidgets.QFileDialog(self) dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptMode.AcceptOpen) dialog.setFileMode(QtWidgets.QFileDialog.FileMode.ExistingFile) - dialog.setNameFilters(valid_files.keys()) + dialog.setNameFilters(valid_loaders.keys()) # dialog.setOption(QtWidgets.QFileDialog.Option.DontUseNativeDialog) if dialog.exec(): files = dialog.selectedFiles() - fn, kargs = valid_files[dialog.selectedNameFilter()] + fn, kargs = valid_loaders[dialog.selectedNameFilter()] # !TODO: handle ambiguous datasets self.slicer_area.set_data(fn(files[0], **kargs)) self.slicer_area.view_all() @@ -449,14 +449,18 @@ def _export_file(self): dialog = QtWidgets.QFileDialog(self) dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptMode.AcceptSave) dialog.setFileMode(QtWidgets.QFileDialog.FileMode.AnyFile) - dialog.setNameFilter( - "xarray HDF5 Files (*.h5)", - ) + + valid_savers: dict[str, tuple[Callable, dict]] = { + "xarray HDF5 Files (*.h5)": (erlab.io.save_as_hdf5, {}), + "NetCDF Files (*.nc *.nc4 *.cdf)": (erlab.io.save_as_netcdf, {}), + } + dialog.setNameFilters(valid_savers.keys()) dialog.setDirectory(f"{self.slicer_area._data.name}.h5") # dialog.setOption(QtWidgets.QFileDialog.Option.DontUseNativeDialog) if dialog.exec(): files = dialog.selectedFiles() - erlab.io.save_as_hdf5(self.slicer_area._data, files[0]) + fn, kargs = valid_savers[dialog.selectedNameFilter()] + fn(self.slicer_area._data, files[0], **kargs) if __name__ == "__main__": @@ -505,9 +509,9 @@ def _export_file(self): data = xr.load_dataarray( # # "~/Documents/ERLab/TiSe2/kxy10.nc", # # "~/Documents/ERLab/TiSe2/221213_SSRL_BL5-2/fullmap_kconv_.h5", - # # "~/Documents/ERLab/CsV3Sb5/2021_Dec_ALS_CV3Sb5/Data/cvs_kxy_small.nc", - # "~/Documents/ERLab/CsV3Sb5/2021_Dec_ALS_CV3Sb5/Data/cvs_kxy.nc", - "~/Documents/ERLab/TiSe2/220410_ALS_BL4/map_mm_4d_.nc", + "~/Documents/ERLab/CsV3Sb5/2021_Dec_ALS_CV3Sb5/Data/cvs_kxy_small.nc", + # "~/Documents/ERLab/CsV3Sb5/2021_Dec_ALS_CV3Sb5/Data/cvs_kxy.nc", + # "~/Documents/ERLab/TiSe2/220410_ALS_BL4/map_mm_4d_.nc", # engine="h5netcdf", ) diff --git a/src/erlab/io/__init__.py b/src/erlab/io/__init__.py index cc547360..601b33c9 100644 --- a/src/erlab/io/__init__.py +++ b/src/erlab/io/__init__.py @@ -58,14 +58,7 @@ import warnings -from erlab.io.dataloader import ( - load, - loader_context, - loaders, - set_data_dir, - set_loader, - summarize, -) +from erlab.io.dataloader import LoaderRegistry from erlab.io.igor import load_experiment, load_wave from erlab.io.utilities import load_hdf5, open_hdf5, save_as_hdf5, save_as_netcdf @@ -73,6 +66,19 @@ # isort: off import erlab.io.plugins # noqa: F401 +loaders = LoaderRegistry.instance() +""" +Global instance of :class:`LoaderRegistry `. + +:meta hide-value: +""" + +load = loaders.load +loader_context = loaders.loader_context +set_data_dir = loaders.set_data_dir +set_loader = loaders.set_loader +summarize = loaders.summarize + merlin = loaders["merlin"] ssrl52 = loaders["ssrl52"] diff --git a/src/erlab/io/dataloader.py b/src/erlab/io/dataloader.py index 5ce605e4..b57474de 100644 --- a/src/erlab/io/dataloader.py +++ b/src/erlab/io/dataloader.py @@ -936,13 +936,3 @@ def _repr_html_(self) -> str: load.__doc__ = LoaderBase.load.__doc__ summarize.__doc__ = LoaderBase.summarize.__doc__ - - -loaders: LoaderRegistry = LoaderRegistry.instance() -"""Global instance of `LoaderRegistry`\n\n:meta hide-value:""" - -set_loader = loaders.set_loader -loader_context = loaders.loader_context -set_data_dir = loaders.set_data_dir -load = loaders.load -summarize = loaders.summarize diff --git a/src/erlab/io/plugins/da30.py b/src/erlab/io/plugins/da30.py index 0af2aa61..96d0a3e8 100644 --- a/src/erlab/io/plugins/da30.py +++ b/src/erlab/io/plugins/da30.py @@ -16,8 +16,8 @@ class DA30Loader(LoaderBase): - name: str = "_da30" - aliases: list[str] = [] + name: str = "da30" + aliases: list[str] = ["DA30"] name_map: dict[str, str] = { "eV": ["Kinetic Energy [eV]", "Energy [eV]"], diff --git a/src/erlab/io/utilities.py b/src/erlab/io/utilities.py index c0fa3902..5b75963f 100644 --- a/src/erlab/io/utilities.py +++ b/src/erlab/io/utilities.py @@ -111,7 +111,7 @@ def fix_attr_format(da: xr.DataArray): return da -def open_hdf5(filename: str | os.PathLike) -> xr.DataArray | xr.Dataset: +def open_hdf5(filename: str | os.PathLike, **kwargs: dict) -> xr.DataArray | xr.Dataset: """Open data from an HDF5 file saved with `save_as_hdf5`. This is a thin wrapper around `xarray.open_dataarray` and `xarray.open_dataset`. @@ -120,19 +120,22 @@ def open_hdf5(filename: str | os.PathLike) -> xr.DataArray | xr.Dataset: ---------- filename The path to the HDF5 file. + **kwargs + Extra arguments to `xarray.open_dataarray` or `xarray.open_dataset`. Returns ------- xarray.DataArray or xarray.Dataset The opened data. """ + kwargs.setdefault("engine", "h5netcdf") try: - return xr.open_dataarray(filename, engine="h5netcdf") + return xr.open_dataarray(filename, **kwargs) except ValueError: - return xr.open_dataset(filename, engine="h5netcdf") + return xr.open_dataset(filename, **kwargs) -def load_hdf5(filename: str | os.PathLike) -> xr.DataArray | xr.Dataset: +def load_hdf5(filename: str | os.PathLike, **kwargs: dict) -> xr.DataArray | xr.Dataset: """Load data from an HDF5 file saved with `save_as_hdf5`. This is a thin wrapper around `xarray.load_dataarray` and `xarray.load_dataset`. @@ -141,16 +144,19 @@ def load_hdf5(filename: str | os.PathLike) -> xr.DataArray | xr.Dataset: ---------- filename The path to the HDF5 file. + **kwargs + Extra arguments to `xarray.load_dataarray` or `xarray.load_dataset`. Returns ------- xarray.DataArray or xarray.Dataset The loaded data. """ + kwargs.setdefault("engine", "h5netcdf") try: - return xr.load_dataarray(filename, engine="h5netcdf") + return xr.load_dataarray(filename, **kwargs) except ValueError: - return xr.load_dataset(filename, engine="h5netcdf") + return xr.load_dataset(filename, **kwargs) def save_as_hdf5( @@ -227,7 +233,6 @@ def save_as_netcdf(data: xr.DataArray, filename: str | os.PathLike, **kwargs: di documentation for a list of all possible arguments. """ - # data = data.assign_attrs(provenance="") kwargs.setdefault("engine", "h5netcdf") fix_attr_format(data).to_netcdf( filename,