Merge branch 'daredevil' of https://github.com/arafune/arpes into dar…

…edevil

arpes/bootstrap.py

@@ -242,7 +242,7 @@ def from_param(cls: type, model_param: lf.Model.Parameter):
         return cls(center=model_param.value, stderr=model_param.stderr)
 
 
-def propagate_errors(f) -> Callable:
+def propagate_errors(f: Callable) -> Callable:
     """A decorator which provides transparent propagation of statistical errors.
 
     The way that this is accommodated is that the inner function is turned into one which
@@ -332,7 +332,12 @@ def bootstrap(
     elif resample_method == "cycle":
         resample_fn = resample_cycle
 
-    def bootstrapped(*args, n: int = 20, prior_adjustment=1, **kwargs: Incomplete):
+    def bootstrapped(
+        *args,
+        n: int = 20,
+        prior_adjustment: int = 1,
+        **kwargs: Incomplete,
+    ):
         # examine args to determine which to resample
         resample_indices = [
             i

arpes/io.py

@@ -172,20 +172,14 @@ def stitch(
     Returns:
         The concatenated data.
     """
-    list_of_files = None
-    if isinstance(df_or_list, pd.DataFrame):
-        list_of_files = list(df_or_list.index)
-    else:
-        if not isinstance(df_or_list, list | tuple):
-            msg = "Expected an interable for a list of the scans to stitch together"
-            raise TypeError(msg)
-        list_of_files = list(df_or_list)
+    list_of_files = _df_or_list_to_files(df_or_list)
     if not built_axis_name:
         assert isinstance(attr_or_axis, str)
         built_axis_name = attr_or_axis
     if not list_of_files:
         msg = "Must supply at least one file to stitch"
         raise ValueError(msg)
+    #
     loaded = [
         f if isinstance(f, xr.DataArray | xr.Dataset) else load_data(f) for f in list_of_files
     ]
@@ -221,6 +215,26 @@ def stitch(
     return concatenated
 
 
+def _df_or_list_to_files(
+    df_or_list: list[str] | pd.DataFrame,
+) -> list[str]:
+    """Helper function for stitch.
+
+    Args:
+        df_or_list(pd.DataFrame, list): input data file
+
+    Returns: (list[str])
+        list of files to stitch.
+    """
+    if isinstance(df_or_list, pd.DataFrame):
+        return list(df_or_list.index)
+    assert not isinstance(
+        df_or_list,
+        list | tuple,
+    ), "Expected an interable for a list of the scans to stitch together"
+    return list(df_or_list)
+
+
 def file_for_pickle(name: str) -> Path | str:
     here = Path()
     from arpes.config import CONFIG

arpes/models/band.py

@@ -1,4 +1,4 @@
-"""Rudimentary band analyis code."""
+"""Rudimentary band analysis code."""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

arpes/plotting/path_tool.py

@@ -134,14 +134,14 @@ def click_main_image(event):
             ("Path", "path"),
         ]
 
-        def convert_to_xarray():
+        def convert_to_xarray() -> dict[str, xr.dataset]:
             """Creates a Dataset consisting of one array for each path.
 
             For each of the paths, we will create a dataset which has an index dimension,
             and datavariables for each of the coordinate dimensions
             """
 
-            def convert_single_path_to_xarray(points):
+            def convert_single_path_to_xarray(points) -> xr.Dataset:
                 vars = {d: np.array([p[i] for p in points]) for i, d in enumerate(self.arr.dims)}
                 coords = {
                     "index": np.array(range(len(points))),

arpes/preparation/axis_preparation.py

@@ -182,8 +182,7 @@ def normalize(arr: xr.Dataset | xr.DataArray) -> xr.DataArray:
 
 def transform_dataarray_axis(
     func: Callable[..., ...],
-    old_axis_name: str,
-    new_axis_name: str,
+    old_and_new_axis_names: tuple[str, str],
     new_axis: NDArray[np.float_] | xr.DataArray,
     dataset: xr.Dataset,
     prep_name: Callable[[str], str],
@@ -195,14 +194,15 @@ def transform_dataarray_axis(
 
     Args:
         func ([TODO:type]): [TODO:description]
-        old_axis_name(str): [TODO:description]
-        new_axis_name(str): [TODO:description]
+        old_and_new_axis_names (tuple[str, str]) : old and new axis names as the tuple form
         new_axis ([TODO:type]): [TODO:description]
         dataset(xr.Dataset): [TODO:description]
         prep_name ([TODO:type]): [TODO:description]
         transform_spectra ([TODO:type]): [TODO:description]
         remove_old ([TODO:type]): [TODO:description]
     """
+    old_axis_name, new_axis_name = old_and_new_axis_names
+
     ds = dataset.copy()
     if transform_spectra is None:
         # transform *all* DataArrays in the dataset that have old_axis_name in their dimensions.

arpes/preparation/coord_preparation.py

@@ -8,15 +8,17 @@
 import numpy as np
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import xarray as xr
 
 __all__ = ["disambiguate_coordinates"]
 
 
 def disambiguate_coordinates(
     datasets: xr.Dataset,
-    possibly_clashing_coordinates,
-):
+    possibly_clashing_coordinates: Sequence[str],
+) -> list[xr.DataArray]:
     """Finds and unifies duplicated coordinates or ambiguous coordinates.
 
     This is useful if two regions claim to have an energy axis, but one is a core level

arpes/preparation/tof_preparation.py

@@ -245,8 +245,7 @@ def process_SToF(dataset: xr.Dataset) -> xr.Dataset:
 
     dataset = transform_dataarray_axis(
         build_KE_coords_to_time_coords(dataset, ke_axis),
-        "time",
-        "eV",
+        ("time", "eV"),
         ke_axis,
         dataset,
         lambda x: x,
@@ -276,8 +275,7 @@ def process_DLD(dataset: xr.Dataset) -> xr.Dataset:
     )
     return transform_dataarray_axis(
         build_KE_coords_to_time_pixel_coords(dataset, ke_axis),
-        "t_pixels",
-        "kinetic",
+        ("t_pixels", "kinetic"),
         ke_axis,
         dataset,
         lambda: "kinetic_spectrum",

arpes/utilities/bz.py

@@ -242,11 +242,6 @@ def flat_bz_indices_list(
         bz_indices_list = [(0, 0)]
 
     assert len(bz_indices_list[0]) in {2, 3}
-    try:
-        if len(bz_indices_list[0]) not in {2, 3}:
-            raise ValueError
-    except (ValueError, TypeError):
-        bz_indices_list = [bz_indices_list]
 
     indices = []
     if len(bz_indices_list[0]) == 2:  # noqa: PLR2004

arpes/utilities/geometry.py

@@ -73,7 +73,12 @@ def segment_contains_point(
     return 0 - epsilon < delta.dot(delta_p) / delta.dot(delta) < 1 + epsilon
 
 
-def polyhedron_intersect_plane(poly_faces, plane_normal, plane_point, epsilon: float = 1e-6):
+def polyhedron_intersect_plane(
+    poly_faces: list[NDArray[np.float_]],
+    plane_normal: NDArray[np.float_],
+    plane_point: NDArray[np.float_],
+    epsilon: float = 1e-6,
+) -> ConvexHull:
     """Determines the intersection of a convex polyhedron intersecting a plane.
 
     The polyhedron faces should be given by a list of np.arrays, where each np.array at
@@ -94,7 +99,7 @@ def polyhedron_intersect_plane(poly_faces, plane_normal, plane_point, epsilon: f
     """
     collected_points = []
 
-    def add_point(c):
+    def add_point(c: NDArray[np.float_]) -> None:
         already_collected = False
         for other in collected_points:
             delta = c - other
@@ -106,7 +111,9 @@ def add_point(c):
             collected_points.append(c)
 
     for poly_face in poly_faces:
-        segments = list(zip(poly_face, np.concatenate([poly_face[1:], [poly_face[0]]])))
+        segments = list(
+            zip(poly_face, np.concatenate([poly_face[1:], [poly_face[0]]]), strict=True),
+        )
         for a, b in segments:
             intersection = point_plane_intersection(
                 plane_normal,

arpes/widgets.py

@@ -29,6 +29,7 @@
 
 import itertools
 import pathlib
+import pprint
 import warnings
 from collections.abc import Sequence
 from functools import wraps
@@ -38,6 +39,7 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
+import pyperclip
 import xarray as xr
 from matplotlib import gridspec
 from matplotlib.axes import Axes
@@ -340,7 +342,7 @@ def mask_cmap(self) -> Colormap:
         return self._mask_cmap
 
     @property
-    def mask(self):
+    def mask(self):  # noqa: ANN202
         return self._mask
 
     @mask.setter
@@ -468,10 +470,6 @@ def on_add_new_peak(selection) -> None:
     ctx["data"] = data
 
     def on_copy_settings(event: MouseEvent) -> None:
-        import pprint
-
-        import pyperclip
-
         pyperclip.copy(pprint.pformat(compute_parameters()))
 
     copy_settings_button = Button(ax_test, "Copy Settings")
@@ -489,7 +487,7 @@ def pca_explorer(
     *,
     transpose_mask: bool = False,
 ) -> CURRENTCONTEXT:
-    """A tool providing PCA decomposition exploration of a dataset.
+    """A tool providing PCA (Principal component analysis) decomposition exploration of a dataset.
 
     Args:
         pca: The decomposition of the data, the output of an sklearn PCA decomp.
@@ -517,7 +515,7 @@ def pca_explorer(
     }
     arpes.config.CONFIG["CURRENT_CONTEXT"] = context
 
-    def compute_for_scatter():
+    def compute_for_scatter() -> tuple[xr.DataArray | xr.Dataset, int]:
         for_scatter = pca.copy(deep=True).isel(
             **dict([[component_dim, context["selected_components"]]]),
         )
@@ -716,18 +714,7 @@ def compute_offsets() -> dict[str, float]:
         return {k: v.val for k, v in sliders.items()}
 
     def on_copy_settings(event: MouseEvent) -> None:
-        try:
-            import pprint
-
-            import pyperclip
-
-            pyperclip.copy(pprint.pformat(compute_offsets()))
-        except ImportError:
-            pass
-        finally:
-            import pprint
-
-            print(pprint.pformat(compute_offsets()))
+        pyperclip.copy(pprint.pformat(compute_offsets()))
 
     def apply_offsets(event: MouseEvent) -> None:
         for name, offset in compute_offsets().items():

arpes/xarray_extensions.py

@@ -387,7 +387,7 @@ def transpose_to_back(self, dim: str) -> xr.DataArray | xr.Dataset:
 
     def select_around_data(
         self,
-        points: dict[str, float] | xr.Dataset,
+        points: dict[str, xr.DataArray] | xr.Dataset,
         radius: dict[str, float] | float | None = None,  # radius={"phi": 0.005}
         *,
         mode: Literal["sum", "mean"] = "sum",
@@ -432,24 +432,9 @@ def select_around_data(
         if isinstance(points, xr.Dataset):
             points = {k: points[k].item() for k in points.data_vars}
 
-        if isinstance(radius, float):
-            radius = {str(d): radius for d in points}
-        else:
-            collected_terms = {f"{k}_r" for k in points}.intersection(
-                set(kwargs.keys()),
-            )
-            if collected_terms:
-                radius = {
-                    str(d): kwargs.get(f"{d}_r", DEFAULT_RADII.get(str(d), UNSPESIFIED))
-                    for d in points
-                }
-            elif radius is None:
-                radius = {str(d): DEFAULT_RADII.get(str(d), UNSPESIFIED) for d in points}
+        radius = self._radius(points, radius, **kwargs)
 
-        assert isinstance(radius, dict)
-        radius = {
-            str(d): radius.get(str(d), DEFAULT_RADII.get(str(d), UNSPESIFIED)) for d in points
-        }
+        logger.debug(f"iter(points.values()): {iter(points.values())}")
 
         along_dims = next(iter(points.values())).dims
         selected_dims = list(points.keys())
@@ -538,25 +523,7 @@ def select_around(
         if isinstance(points, xr.Dataset):
             points = {k: points[k].item() for k in points.data_vars}
         logger.debug(f"points: {points}")
-        if isinstance(radius, float):
-            radius = {str(d): radius for d in points}
-        else:
-            collected_terms = {f"{k}_r" for k in points}.intersection(
-                set(kwargs.keys()),
-            )
-            if collected_terms:
-                radius = {
-                    str(d): kwargs.get(f"{d}_r", DEFAULT_RADII.Get(str(d), UNSPESIFIED))
-                    for d in points
-                }
-            elif radius is None:
-                radius = {str(d): DEFAULT_RADII.get(str(d), UNSPESIFIED) for d in points}
-
-        assert isinstance(radius, dict)
-        radius = {
-            str(d): radius.get(str(d), DEFAULT_RADII.get(str(d), UNSPESIFIED)) for d in points
-        }
-
+        radius = self._radius(points, radius, **kwargs)
         logger.debug(f"radius: {radius}")
         nearest_sel_params = {}
 
@@ -585,6 +552,38 @@ def select_around(
             return selected.sum(list(radius.keys()))
         return selected.mean(list(radius.keys()))
 
+    def _radius(
+        self,
+        points: dict[str, float],
+        radius: float | dict[str, float] | None,
+        **kwargs: float,
+    ) -> dict[str, float]:
+        """Helper function. Generate radius dict.
+
+        When radius is dict form, nothing has been done, essentially.
+
+        Args:
+            points (dict[str, float]): Selection point
+            radius (dict[str, float] | float | None): radius
+            kwargs (float): [TODO:description]
+
+        Returns: dict[str, float]
+            radius for selection.
+        """
+        if isinstance(radius, float):
+            radius = {str(d): radius for d in points}
+        else:
+            collectted_terms = {f"{k}_r" for k in points}.intersection(set(kwargs.keys()))
+            if collectted_terms:
+                radius = {
+                    str(d): kwargs.get(f"{d}_r", DEFAULT_RADII.get(str(d), UNSPESIFIED))
+                    for d in points
+                }
+            elif radius is None:
+                radius = {str(d): DEFAULT_RADII.get(str(d), UNSPESIFIED) for d in points}
+        assert isinstance(radius, dict)
+        return {str(d): radius.get(str(d), DEFAULT_RADII.get(str(d), UNSPESIFIED)) for d in points}
+
     def short_history(self, key: str = "by") -> list:
         """Return the short version of history.