From e3220e8f12f3b6fb4e77139ef921ee336d3f2473 Mon Sep 17 00:00:00 2001
From: Ryuichi Arafune <ryuichi.arafune@gmail.com>
Date: Thu, 1 Feb 2024 12:28:17 +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

---
 arpes/analysis/band_analysis.py | 14 +++++++++-----
 arpes/analysis/deconvolution.py | 13 +++++--------
 arpes/analysis/mask.py          |  1 +
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/arpes/analysis/band_analysis.py b/arpes/analysis/band_analysis.py
index 364ae127..c3e808d0 100644
--- a/arpes/analysis/band_analysis.py
+++ b/arpes/analysis/band_analysis.py
@@ -1,4 +1,5 @@
 """Provides some band analysis tools."""
+
 from __future__ import annotations
 
 import contextlib
@@ -21,6 +22,7 @@
 from arpes.utilities.conversion.forward import convert_coordinates_to_kspace_forward
 from arpes.utilities.jupyter import wrap_tqdm
 
+
 if TYPE_CHECKING:
     from collections.abc import Generator
 
@@ -36,7 +38,10 @@
 )
 
 
-def fit_for_effective_mass(data: DataType, fit_kwargs: dict | None = None) -> float:
+def fit_for_effective_mass(
+    data: DataType,
+    fit_kwargs: dict | None = None,
+) -> float:
     """Fits for the effective mass in a piece of data.
 
     Performs an effective mass fit by first fitting for Lorentzian lineshapes and then fitting
@@ -77,7 +82,7 @@ def fit_for_effective_mass(data: DataType, fit_kwargs: dict | None = None) -> fl
         final_mom = next(dim for dim in ["kx", "ky", "kp", "kz"] if dim in forward)
         eVs = results.F.p("a_center").values
         kps = [
-            forward[final_mom].sel(eV=eV, **dict([[mom_dim, ang]]), method="nearest")
+            forward[final_mom].sel(dict([[mom_dim, ang]]), eV=eV, method="nearest")
             for eV, ang in zip(eVs, data_array.coords[mom_dim].values, strict=True)
         ]
         quad_fit = QuadraticModel().fit(eVs, x=np.array(kps))
@@ -90,7 +95,7 @@ def fit_for_effective_mass(data: DataType, fit_kwargs: dict | None = None) -> fl
 
 def unpack_bands_from_fit(
     band_results: xr.DataArray,
-    weights: tuple[float, float, float] = tuple(),
+    weights: tuple[float, float, float] = (2, 0, 10),
 ) -> list[arpes.models.band.Band]:
     """Deconvolve the band identities of a series of overlapping bands.
 
@@ -124,8 +129,6 @@ def unpack_bands_from_fit(
     Returns:
         Unpacked bands.
     """
-    if not weights:
-        weights = (2, 0, 10)
 
     template_components = band_results.values[0].model.components
     prefixes = [component.prefix for component in template_components]
@@ -290,6 +293,7 @@ def fit_patterned_bands(
         Dataset or DataArray, as controlled by the parameter "dataset"
     """
     if background:
+
         from arpes.models.band import AffineBackgroundBand
 
         background = AffineBackgroundBand
diff --git a/arpes/analysis/deconvolution.py b/arpes/analysis/deconvolution.py
index f3ab0209..af2b8066 100644
--- a/arpes/analysis/deconvolution.py
+++ b/arpes/analysis/deconvolution.py
@@ -1,4 +1,5 @@
 """Provides deconvolution implementations, especially for 2D Richardson-Lucy."""
+
 from __future__ import annotations
 
 import contextlib
@@ -46,7 +47,7 @@ def deconvolve_ice(
     The PSF is the impulse response of a focused optical imaging system.
 
     Args:
-        data: input data
+        data (DataType): input data
         psf(NDArray[np.float_): array as point spread function
         n_iterations: the number of convolutions to use for the fit
         deg: the degree of the fitting polynominial
@@ -54,11 +55,7 @@ def deconvolve_ice(
     Returns:
         The deconvoled data in the same format.
     """
-    arr = normalize_to_spectrum(data)
-    if type(data) is np.ndarray:
-        pass
-    else:
-        arr = arr.values
+    arr = normalize_to_spectrum(data).values
 
     if deg is None:
         deg = n_iterations - 3
@@ -100,8 +97,8 @@ def deconvolve_rl(
     Args:
         data: input data
         axis
-        sigma
         mode: pass to ndimage.convolve
+        sigma
         progress
         psf: for 1d, if not specified, must specify axis and sigma
         n_iterations: the number of convolutions to use for the fit
@@ -235,7 +232,7 @@ def wrap_progress(
                 # need to explicitly specify for some versions of numpy
 
             result = u[-1]
-    else:
+    else:  # data.dims == 1
         if type(arr) is not np.ndarray:
             arr = arr.values
         u = [arr]
diff --git a/arpes/analysis/mask.py b/arpes/analysis/mask.py
index 8ac2e24e..94ec6ab6 100644
--- a/arpes/analysis/mask.py
+++ b/arpes/analysis/mask.py
@@ -1,4 +1,5 @@
 """Utilities for applying masks to data."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING