From 36ad17f1bb0493c44b762ec3f0dd11ee2495aeb1 Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Wed, 18 Oct 2023 11:27:04 +0900 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A8=20=20update=20type=20hints?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../plugin/Elettra_spectromicroscopy.py | 2 +- arpes/io.py | 17 +++++- arpes/utilities/conversion/forward.py | 59 ++++++++++++------- arpes/xarray_extensions.py | 12 ++-- pyproject.toml | 2 +- 5 files changed, 61 insertions(+), 31 deletions(-) diff --git a/arpes/endstations/plugin/Elettra_spectromicroscopy.py b/arpes/endstations/plugin/Elettra_spectromicroscopy.py index 1f3f900e..325fc1d5 100644 --- a/arpes/endstations/plugin/Elettra_spectromicroscopy.py +++ b/arpes/endstations/plugin/Elettra_spectromicroscopy.py @@ -27,7 +27,7 @@ def collect_coord(index: int, dset: h5py.Dataset) -> tuple[str, NDArray[np.float Args: index: The index of the coordinate to extract from metadata. - dset: The HDF dataset containin Elettra spectromicroscopy data. + dset: The HDF dataset containing Elettra spectromicroscopy data. Returns: The coordinate extracted at `index` from the metadata. The return convention here is to diff --git a/arpes/io.py b/arpes/io.py index 66da07d7..f94db6d1 100644 --- a/arpes/io.py +++ b/arpes/io.py @@ -16,6 +16,7 @@ import warnings from collections.abc import Iterable from dataclasses import dataclass +from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger from pathlib import Path from typing import TYPE_CHECKING @@ -37,6 +38,19 @@ ) +LOGLEVELS = (DEBUG, INFO) +LOGLEVEL = LOGLEVELS[0] +logger = getLogger(__name__) +fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s" +formatter = Formatter(fmt) +handler = StreamHandler() +handler.setLevel(LOGLEVEL) +logger.setLevel(LOGLEVEL) +handler.setFormatter(formatter) +logger.addHandler(handler) +logger.propagate = False + + def load_data( file: str | Path | int, location: str | type | None = None, @@ -84,7 +98,8 @@ def load_data( ), stacklevel=2, ) - + if kwargs.get("trace"): + logger.debug(f"contents of desc: {desc}") return load_scan(desc, **kwargs) diff --git a/arpes/utilities/conversion/forward.py b/arpes/utilities/conversion/forward.py index 8a864bd9..1ecb53e1 100644 --- a/arpes/utilities/conversion/forward.py +++ b/arpes/utilities/conversion/forward.py @@ -12,6 +12,7 @@ from __future__ import annotations import warnings +from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger from typing import TYPE_CHECKING import numpy as np @@ -19,7 +20,7 @@ from arpes.analysis.filters import gaussian_filter_arr from arpes.provenance import update_provenance -from arpes.trace import traceable +from arpes.trace import Trace, traceable from arpes.utilities import normalize_to_spectrum from .bounds_calculations import ( @@ -31,7 +32,7 @@ from .core import convert_to_kspace if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Sequence from numpy.typing import NDArray @@ -46,12 +47,25 @@ ) +LOGLEVELS = (DEBUG, INFO) +LOGLEVEL = LOGLEVELS[1] +logger = getLogger(__name__) +fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s" +formatter = Formatter(fmt) +handler = StreamHandler() +handler.setLevel(LOGLEVEL) +logger.setLevel(LOGLEVEL) +handler.setFormatter(formatter) +logger.addHandler(handler) +logger.propagate = False + + @traceable def convert_coordinate_forward( data: DataType, coords: dict[str, float], *, - trace: Callable = None, # noqa: RUF013 + trace: Trace | None = None, **k_coords: NDArray[np.float_], ) -> dict[str, float]: """Inverse/forward transform for the small angle volumetric k-conversion code. @@ -83,8 +97,8 @@ def convert_coordinate_forward( Another approach would be to write down the exact small angle approximated transforms. Args: - data: The data defining the coordinate offsets and experiment geometry. - coords: The coordinates of a point in angle-space to be converted. + data (DataType): The data defining the coordinate offsets and experiment geometry. + coords (dict[str, float]): The coordinates of a *point* in angle-space to be converted. trace: Used for performance tracing and debugging. k_coords: Coordinate for k-axis @@ -110,20 +124,20 @@ def convert_coordinate_forward( "ky": np.linspace(-4, 4, 300), } # Copying after taking a constant energy plane is much much cheaper - trace("Copying") + trace("Copying") if trace else None data_arr = data_arr.copy(deep=True) data_arr.loc[data_arr.G.round_coordinates(coords)] = data_arr.values.max() * 100000 - trace("Filtering") + trace("Filtering") if trace else None data_arr = gaussian_filter_arr(data_arr, default_size=3) - trace("Converting once") + trace("Converting once") if trace else None kdata = convert_to_kspace(data_arr, **k_coords, trace=trace) - trace("argmax") + trace("argmax") if trace else None near_target = kdata.G.argmax_coords() - trace("Converting twice") + trace("Converting twice") if trace else None kdata_close = convert_to_kspace( data_arr, trace=trace, @@ -132,7 +146,7 @@ def convert_coordinate_forward( # inconsistently, the energy coordinate is sometimes returned here # so we remove it just in case - trace("argmax") + trace("argmax") if trace else None coords = kdata_close.G.argmax_coords() if "eV" in coords: del coords["eV"] @@ -148,7 +162,7 @@ def convert_through_angular_pair( # noqa: PLR0913 transverse_specification: dict[str, NDArray[np.float_]], *, relative_coords: bool = True, - trace: Callable = None, # noqa: RUF013 + trace: Trace | None = None, **k_coords: NDArray[np.float_], ) -> dict[str, float]: """Converts the lower dimensional ARPES cut passing through `first_point` and `second_point`. @@ -203,12 +217,12 @@ def convert_through_angular_pair( # noqa: PLR0913 k_second_point["ky"] - k_first_point["ky"], k_second_point["kx"] - k_first_point["kx"], ) - trace(f"Determined offset angle {-offset_ang}") + trace(f"Determined offset angle {-offset_ang}") if trace else None with data.S.with_rotation_offset(-offset_ang): - trace("Finding first momentum coordinate.") + trace("Finding first momentum coordinate.") if trace else None k_first_point = convert_coordinate_forward(data, first_point, trace=trace, **k_coords) - trace("Finding second momentum coordinate.") + trace("Finding second momentum coordinate.") if trace else None k_second_point = convert_coordinate_forward(data, second_point, trace=trace, **k_coords) # adjust output coordinate ranges @@ -229,7 +243,7 @@ def convert_through_angular_pair( # noqa: PLR0913 parallel_axis = np.linspace(left_point, right_point, len(parallel_axis)) # perform the conversion - trace("Performing final momentum conversion.") + trace("Performing final momentum conversion.") if trace else None converted_data = convert_to_kspace( data, **transverse_specification, @@ -237,7 +251,7 @@ def convert_through_angular_pair( # noqa: PLR0913 trace=trace, ).mean(list(transverse_specification.keys())) - trace("Annotating the requested point momentum values.") + trace("Annotating the requested point momentum values.") if trace else None return converted_data.assign_attrs( { "first_point_kx": k_first_point[parallel_dim], @@ -255,7 +269,7 @@ def convert_through_angular_point( # noqa: PLR0913 transverse_specification: dict[str, NDArray[np.float_]], *, relative_coords: bool = True, - trace: Callable = None, # noqa: RUF013 + trace: Trace | None = None, **k_coords: NDArray[np.float_], ) -> xr.DataArray: """Converts the lower dimensional ARPES cut passing through given angular `coords`. @@ -311,7 +325,7 @@ def convert_coordinates( ) -> xr.Dataset: """Converts coordinates forward in momentum.""" - def unwrap_coord(coord): + def unwrap_coord(coord: xr.DataArray | NDArray[np.float_]) -> NDArray[np.float_]: try: return coord.values except (TypeError, AttributeError): @@ -389,7 +403,7 @@ def expand_to(cname: str, c: Sequence[float]) -> float: @update_provenance("Forward convert coordinates to momentum") -def convert_coordinates_to_kspace_forward(arr: DataType) -> xr.Dataset | None: +def convert_coordinates_to_kspace_forward(arr: DataType) -> xr.Dataset: """Forward converts all the individual coordinates of the data array. Args: @@ -406,7 +420,8 @@ def convert_coordinates_to_kspace_forward(arr: DataType) -> xr.Dataset | None: momentum_compatibles: list[str] = list(all_indexes.keys()) momentum_compatibles.sort() if not momentum_compatibles: - return None + msg = "Cannot convert because no momentum compatible coordinate" + raise RuntimeError(msg) dest_coords = { ("phi",): ["kp", "kz"], ("theta",): ["kp", "kz"], @@ -443,7 +458,7 @@ def broadcast_by_dim_location( # else we are dealing with an actual array the_slice = [None] * len(target_shape) the_slice[dim_location] = slice(None, None, None) - print(dim_location) + logger.info(dim_location) return np.asarray(data)[the_slice] raw_coords = { diff --git a/arpes/xarray_extensions.py b/arpes/xarray_extensions.py index 198b5383..39b87901 100644 --- a/arpes/xarray_extensions.py +++ b/arpes/xarray_extensions.py @@ -46,7 +46,7 @@ from collections import OrderedDict from collections.abc import Collection, Sequence from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger -from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias, Unpack +from typing import TYPE_CHECKING, Any, Literal, NoReturn, Self, TypeAlias, Unpack import matplotlib.pyplot as plt import numpy as np @@ -316,7 +316,7 @@ def hv(self) -> float | None: return 5.93 return None - def fetch_ref_attrs(self) -> Incomplete: + def fetch_ref_attrs(self) -> dict[str, Any]: """Get reference attrs.""" if "ref_attrs" in self._obj.attrs: return self._obj.attrs @@ -3370,7 +3370,7 @@ def is_spatial(self) -> bool: return self.spectrum.S.is_spatial @property - def spectrum(self) -> xr.DataArray | None: + def spectrum(self) -> xr.DataArray | NoReturn: """Isolates a single spectrum from a dataset. This is a convenience method which is typically used in startup for @@ -3577,15 +3577,15 @@ def energy_notation(self) -> EnergyNotation: .. Note:: The "Kinetic" energy refers to the Fermi level. (not Vacuum level) """ if "energy_notation" in self._obj.attrs: - if self.S.spectrum.attrs["energy_notation"] in ( + if self.spectrum.attrs["energy_notation"] in ( "Kinetic", "kinetic", "kinetic energy", ): - self.S.spectrum.attrs["energy_notation"] = "Kinetic" + self.spectrum.attrs["energy_notation"] = "Kinetic" return "Kinetic" return "Binding" - self.S.spectrum.attrs["energy_notation"] = self.S.spectrum.attrs.get( + self.spectrum.attrs["energy_notation"] = self.spectrum.attrs.get( "energy_notation", "Binding", ) diff --git a/pyproject.toml b/pyproject.toml index 07901edc..911eb8cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ ignore = [ "N802", # invalid-function-name (N802) "N806", # non-lowercase-variable-in-function "N999", # invalid-module-name (N999) - "S101", + "S101", # assert (S101) "TD002", # missing-todo-author "TD003", # missing-todo-link "PD011", # pandas-use-of-dot-values