From 3e35287004c4f57689ee7b55167b28aefd62bb22 Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Tue, 27 Feb 2024 10:35:38 +0900 Subject: [PATCH] =?UTF-8?q?=F0=9F=92=AC=20=20Update=20type=20hints?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/source/conf.py | 8 +++++- src/arpes/plotting/false_color.py | 7 +++--- src/arpes/plotting/fit_tool/__init__.py | 25 +++++++++---------- .../plotting/fit_tool/fit_inspection_plot.py | 2 +- src/arpes/plotting/movie.py | 10 ++++---- src/arpes/plotting/qt_ktool/__init__.py | 2 +- src/arpes/plotting/qt_tool/__init__.py | 16 +++++++++--- src/arpes/plotting/utils.py | 10 +++----- 8 files changed, 47 insertions(+), 33 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index b8f7369d..e812767c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -152,7 +152,13 @@ def setup(app): latex_elements = {} latex_documents = [ - (master_doc, "arpes.tex", "arpes Documentation", "Conrad Stansbury", "manual"), + ( + master_doc, + "arpes.tex", + "arpes Documentation", + "Conrad Stansbury/Ryuichi Arafune (>= V4)", + "manual", + ), ] # -- Options for manual page output ------------------------------------------ diff --git a/src/arpes/plotting/false_color.py b/src/arpes/plotting/false_color.py index 36d47eb4..21c048c5 100644 --- a/src/arpes/plotting/false_color.py +++ b/src/arpes/plotting/false_color.py @@ -24,18 +24,19 @@ @save_plot_provenance -def false_color_plot( # noqa: PLR0913 +def false_color_plot( data_rgb: tuple[xr.Dataset, xr.Dataset, xr.Dataset], ax: Axes | None = None, out: str | Path = "", *, invert: bool = False, - pmin: float = 0, - pmax: float = 1, + pmin_pmax: tuple[float, float] = (0, 1), **kwargs: Incomplete, ) -> Path | tuple[Figure | None, Axes]: """Plots a spectrum in false color after conversion to R, G, B arrays.""" data_r_arr, data_g_arr, data_b_arr = (normalize_to_spectrum(d) for d in data_rgb) + pmin, pmax = pmin_pmax + fig: Figure | None = None if ax is None: fig, ax = plt.subplots(figsize=kwargs.pop("figsize", (7, 5))) diff --git a/src/arpes/plotting/fit_tool/__init__.py b/src/arpes/plotting/fit_tool/__init__.py index 22080f45..6d805cc3 100644 --- a/src/arpes/plotting/fit_tool/__init__.py +++ b/src/arpes/plotting/fit_tool/__init__.py @@ -34,6 +34,8 @@ from .fit_inspection_plot import FitInspectionPlot if TYPE_CHECKING: + from collections.abc import Iterable + from _typeshed import Incomplete __all__ = ( @@ -74,13 +76,13 @@ def compile_key_bindings(self) -> list[KeyBinding]: KeyBinding("Transpose - Swap Front Axes", [QtCore.Qt.Key.Key_Y], self.transpose_swap), ] - def center_cursor(self, event) -> None: + def center_cursor(self) -> None: self.app().center_cursor() - def transpose_roll(self, event) -> None: + def transpose_roll(self) -> None: self.app().transpose_to_front(-1) - def transpose_swap(self, event) -> None: + def transpose_swap(self) -> None: self.app().transpose_to_front(1) @staticmethod @@ -93,7 +95,7 @@ def _update_scroll_delta(delta: tuple[int, int], event: QtGui.QKeyEvent) -> tupl return delta - def reset_intensity(self, event: QtGui.QKeyEvent) -> None: + def reset_intensity(self) -> None: self.app().reset_intensity() def scroll_z(self, event: QtGui.QKeyEvent) -> None: @@ -159,7 +161,7 @@ def center_cursor(self) -> None: for cursor in cursors: cursor.set_location(new_cursor[i]) - def scroll(self, delta) -> None: + def scroll(self, delta: Iterable[float]) -> None: """Scroll the axis delta[0] by delta[1] pixels.""" if delta[0] >= len(self.context["cursor"]): warnings.warn("Tried to scroll a non-existent dimension.", stacklevel=2) @@ -233,8 +235,7 @@ def configure_image_widgets(self) -> None: self.generate_marginal_for((), 0, 0, "xy", cursors=True, layout=self.content_layout) self.generate_fit_marginal_for( (0, 1), - 0, - 1, + (0, 1), "fit", cursors=False, orientation=PlotOrientation.Vertical, @@ -246,8 +247,7 @@ def configure_image_widgets(self) -> None: self.generate_marginal_for((2,), 1, 0, "xy", cursors=True, layout=self.content_layout) self.generate_fit_marginal_for( (0, 1, 2), - 0, - 0, + (0, 0), "fit", cursors=True, layout=self.content_layout, @@ -260,8 +260,7 @@ def configure_image_widgets(self) -> None: self.generate_marginal_for((0, 3), 1, 1, "yz", layout=self.content_layout) self.generate_fit_marginal_for( (0, 1, 2, 3), - 0, - 0, + (0, 0), "fit", cursors=True, layout=self.content_layout, @@ -270,8 +269,7 @@ def configure_image_widgets(self) -> None: def generate_fit_marginal_for( self, dimensions: tuple[int, ...], - column: int, - row: int, + column_row: tuple[int, int], name: str = "fit", orientation: PlotOrientation = PlotOrientation.Horizontal, *, @@ -283,6 +281,7 @@ def generate_fit_marginal_for( This does something very similar to `generate_marginal_for` except that it is specialized to showing a widget which embeds information about the current fit result. """ + column, row = column_row if layout is None: layout = self._layout assert isinstance(layout, QLayout) diff --git a/src/arpes/plotting/fit_tool/fit_inspection_plot.py b/src/arpes/plotting/fit_tool/fit_inspection_plot.py index 2eee726b..3a9dc023 100644 --- a/src/arpes/plotting/fit_tool/fit_inspection_plot.py +++ b/src/arpes/plotting/fit_tool/fit_inspection_plot.py @@ -35,7 +35,7 @@ def __init__(self, parent: QWidget | None = None) -> None: def set_model_result(self, model_result: lmfit.model.ModelResult) -> None: """Converts the ModelResult to the HTML representation and sets page contents.""" assert model_result is not None - self.setText(model_result._repr_multiline_text_(short=True)) + self.setText(model_result._repr_multiline_text_(short=True)) # noqa: SLF001 class FitInspectionPlot(QWidget): diff --git a/src/arpes/plotting/movie.py b/src/arpes/plotting/movie.py index 242796c3..3f379121 100644 --- a/src/arpes/plotting/movie.py +++ b/src/arpes/plotting/movie.py @@ -30,7 +30,7 @@ def plot_movie( data: xr.DataArray, time_dim: str = "delay", - interval: float = 100, + interval_ms: float = 100, fig_ax: tuple[Figure | None, Axes | None] = (None, None), out: str | Path = "", **kwargs: Unpack[PColorMeshKwargs], @@ -39,8 +39,8 @@ def plot_movie( Args: data (xr.DataArray): ARPES data - time_dim (str): dimension name for time - interval: [TODO:description] + time_dim (str): dimension name for time, default is "delay". + interval_ms: Delay between frames in milliseconds. fig_ax (tuple[Figure, Axes]): matplotlib object out: [TODO:description] kwargs: [TODO:description] @@ -98,13 +98,13 @@ def animate(i: int) -> tuple[QuadMesh]: init_func=init, repeat=500, frames=len(animation_coords), - interval=interval, + interval=interval_ms, blit=True, ) animation_writer = animation.writers["ffmpeg"] writer = animation_writer( - fps=1000 / interval, + fps=1000 / interval_ms, metadata={"artist": "Me"}, bitrate=1800, ) diff --git a/src/arpes/plotting/qt_ktool/__init__.py b/src/arpes/plotting/qt_ktool/__init__.py index aa7aa55a..5ad1b885 100644 --- a/src/arpes/plotting/qt_ktool/__init__.py +++ b/src/arpes/plotting/qt_ktool/__init__.py @@ -22,7 +22,7 @@ from matplotlib.colors import Colormap from PySide6.QtWidgets import QGridLayout - from arpes._typing import ANGLE, XrTypes + from arpes._typing import ANGLE __all__ = ( "KTool", diff --git a/src/arpes/plotting/qt_tool/__init__.py b/src/arpes/plotting/qt_tool/__init__.py index 9d99e990..f27c32f5 100644 --- a/src/arpes/plotting/qt_tool/__init__.py +++ b/src/arpes/plotting/qt_tool/__init__.py @@ -34,6 +34,8 @@ from .BinningInfoWidget import BinningInfoWidget if TYPE_CHECKING: + from collections.abc import Iterable + from _typeshed import Incomplete from PySide6.QtCore import QEvent from PySide6.QtGui import QKeyEvent @@ -212,7 +214,7 @@ def center_cursor(self) -> None: for cursor in cursors: cursor.set_location(new_cursor[i]) - def scroll(self, delta) -> None: + def scroll(self, delta: Iterable[float]) -> None: """Scroll the axis delta[0] by delta[1] pixels.""" if delta[0] >= len(self.context["cursor"]): warnings.warn("Tried to scroll a non-existent dimension.", stacklevel=2) @@ -236,10 +238,18 @@ def binning(self) -> list[int]: return list(self._binning) @binning.setter - def binning(self, value) -> None: + def binning(self, value: float) -> None: """Set the desired axis binning.""" different_binnings = [ - i for i, (nv, v) in enumerate(zip(value, self._binning, strict=True)) if nv != v + i + for i, (nv, v) in enumerate( + zip( + value, + self._binning, + strict=True, + ), + ) + if nv != v ] self._binning = value diff --git a/src/arpes/plotting/utils.py b/src/arpes/plotting/utils.py index 2be45c6e..2ca316a8 100644 --- a/src/arpes/plotting/utils.py +++ b/src/arpes/plotting/utils.py @@ -30,6 +30,7 @@ from arpes import VERSION from arpes._typing import IMshowParam, XrTypes from arpes.config import CONFIG, SETTINGS, attempt_determine_workspace, is_using_tex +from arpes.constants import TWO_DIMENSION from arpes.utilities import normalize_to_spectrum from arpes.utilities.jupyter import get_notebook_name, get_recent_history @@ -121,9 +122,6 @@ logger.propagate = False -TwoDimensional = 2 - - @contextlib.contextmanager def unchanged_limits(ax: Axes) -> Iterator[None]: """Context manager that retains axis limits.""" @@ -650,7 +648,7 @@ def plot_arr( except AttributeError: n_dims = 1 - if n_dims == TwoDimensional: + if n_dims == TWO_DIMENSION: quad = None if arr is not None: ax, quad = imshow_arr(arr, ax=ax, over=over, **kwargs) @@ -873,7 +871,7 @@ def resolve(name: str, value: slice | int) -> NDArray[np.float_]: assert reference_data is not None logger.info(missing_dims) - if n_cut_dims == TwoDimensional: + if n_cut_dims == TWO_DIMENSION: # a region cut, illustrate with a rect or by suppressing background return @@ -1174,7 +1172,7 @@ def calculate_aspect_ratio(data: xr.DataArray) -> float: """Calculate the aspect ratio which should be used for plotting some data based on extent.""" data_arr = data if isinstance(data, xr.DataArray) else normalize_to_spectrum(data) - assert len(data.dims_arr) == TwoDimensional + assert len(data.dims_arr) == TWO_DIMENSION x_extent = np.ptp(data_arr.coords[data_arr.dims[0]].values) y_extent = np.ptp(data_arr.coords[data_arr.dims[1]].values)