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

…edevil

src/arpes/analysis/__init__.py

@@ -30,5 +30,5 @@
     pocket_parameters,
     radial_edcs_along_pocket,
 )
-from .tarpes import normalized_relative_change, relative_change
+from .tarpes import build_crosscorrelation, normalized_relative_change, relative_change
 from .xps import approximate_core_levels

src/arpes/analysis/decomposition.py

@@ -77,7 +77,7 @@ class FactorAnalysisParam(TypedDict, total=False):
     random_state: int | None
 
 
-class DecompositionParam(PCAParam, FastICAParam, NMFParam, FactorAnalysisParam):
+class DecompositionParam(PCAParam, FastICAParam, NMFParam, FactorAnalysisParam):  # type: ignore[misc]
     pass
 
 

src/arpes/analysis/deconvolution.py

@@ -82,11 +82,8 @@ def deconvolve_ice(
         poly = np.poly1d(coefs)
         deconv[t] = poly(0)
 
-    if isinstance(data, np.ndarray):
-        result = deconv
-    else:
-        result = data.copy(deep=True)
-        result.values = deconv
+    result = data.copy(deep=True)
+    result.values = deconv
     return result
 
 
@@ -114,7 +111,11 @@ def deconvolve_rl(
 
 
 @update_provenance("Make 1D-Point Spread Function")
-def make_psf1d(data: xr.DataArray, dim: str, sigma: float) -> xr.DataArray:
+def make_psf1d(
+    data: xr.DataArray,
+    dim: str,
+    sigma: float,
+) -> xr.DataArray:
     """Produces a 1-dimensional gaussian point spread function for use in deconvolve_rl.
 
     Args:
@@ -184,7 +185,6 @@ def make_psf(
                 f" psf_coords[{k}]: ±{np.max(v):.3f}",
             )
     coords = np.meshgrid(*[psf_coords[dim] for dim in data.dims], indexing="ij")
-
     coords_for_pdf_pos = np.stack(coords, axis=-1)  # point distribution function (pdf)
     logger.debug(f"shape of coords_for_pdf_pos: {coords_for_pdf_pos.shape}")
     psf = xr.DataArray(

src/arpes/analysis/savitzky_golay.py

@@ -3,33 +3,34 @@
 from __future__ import annotations
 
 from math import factorial
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Literal, TypeVar
 
 import numpy as np
 import scipy.signal
 import xarray as xr
+from numpy.typing import NDArray
 
 from arpes.constants import TWO_DIMENSION
 from arpes.provenance import update_provenance
 
 if TYPE_CHECKING:
     from collections.abc import Hashable
 
-    from numpy.typing import NDArray
-
 
 __all__ = ("savitzky_golay",)
 
+T = TypeVar("T", xr.DataArray, NDArray[np.float_])
+
 
 @update_provenance("Savitzky Golay Filter")
 def savitzky_golay(  # noqa: PLR0913
-    data: xr.DataArray,
+    data: T,
     window_size: int,
     order: int,
-    deriv: int | Literal["col", "row", "both", None] = 0,
+    deriv: int | Literal["col", "row", "both"] | None = 0,
     rate: int = 1,
     dim: Hashable = "",
-) -> xr.DataArray:
+) -> T:
     """Implements a Savitzky Golay filter with given window size.
 
     You can specify "pass through" dimensions
@@ -38,6 +39,7 @@ def savitzky_golay(  # noqa: PLR0913
 
     Args:
         data: Input data.
+            This should be xr.DataArray, while list[float] or np.ndarray can be accepted.
         window_size: Number of points in the window that the filter uses locally.
         order: The polynomial order used in the convolution.
         deriv: the order of the derivative to compute (default = 0 means only smoothing)
@@ -47,12 +49,12 @@ def savitzky_golay(  # noqa: PLR0913
     Returns:
         Smoothed data.
     """
-    if isinstance(
-        data,
-        list | np.ndarray,
-    ):
+    if isinstance(data, list):
+        assert isinstance(deriv, int)
+        return savitzky_golay_array(data, window_size, order, deriv, rate)
+    if isinstance(data, np.ndarray):
+        assert isinstance(deriv, int)
         return savitzky_golay_array(data, window_size, order, deriv, rate)
-
     if len(data.dims) == 1:
         assert isinstance(deriv, int)
         transformed_data = savitzky_golay_array(data.values, window_size, order, deriv, rate)

src/arpes/analysis/statistics.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger
 from typing import TYPE_CHECKING
 
 import xarray as xr
@@ -14,10 +15,26 @@
 
 __all__ = ("mean_and_deviation",)
 
+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
+
 
 @update_provenance("Calculate mean and standard deviation for observation axis")
 @lift_dataarray_to_generic
-def mean_and_deviation(data: XrTypes, axis: str = "", name: str = "") -> xr.Dataset:
+def mean_and_deviation(
+    data: xr.DataArray,  # data.name is used.
+    axis: str = "",
+    name: str = "",
+) -> xr.Dataset:
     """Calculates the mean and standard deviation of a DataArray along an axis.
 
     The reduced axis corresponds to individual observations of a tensor/array valued quantity.
@@ -27,7 +44,7 @@ def mean_and_deviation(data: XrTypes, axis: str = "", name: str = "") -> xr.Data
     If a name is not attached to the DataArray, it should be provided.
 
     Args:
-        data: The input data (Both DataArray and Dataset).
+        data: The input data.
         axis: The name of the dimension which we should perform the reduction along.
         name: The name of the variable which should be reduced. By default, uses `data.name`.
 

src/arpes/analysis/tarpes.py

@@ -3,15 +3,109 @@
 from __future__ import annotations
 
 import warnings
+from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger
+from typing import TYPE_CHECKING, TypeVar
 
 import numpy as np
 import xarray as xr
+from numpy.typing import NDArray
 
 from arpes.preparation import normalize_dim
 from arpes.provenance import update_provenance
 from arpes.utilities import normalize_to_spectrum
 
-__all__ = ("find_t0", "relative_change", "normalized_relative_change")
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__all__ = (
+    "find_t0",
+    "relative_change",
+    "normalized_relative_change",
+    "build_crosscorrelation",
+    "delaytime_fs",
+    "position_to_delaytime",
+)
+
+
+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
+
+
+A = TypeVar("A", NDArray[np.float64], float)
+
+
+def build_crosscorrelation(
+    datalist: Sequence[xr.DataArray],
+    delayline_dim: str = "position",
+    delayline_origin: float = 0,
+    *,
+    convert_position_to_time: bool = True,
+) -> xr.DataArray:
+    """Build the ('original dimnsion' + 1)D data from the series of cross-correlation measurements.
+
+    Args:
+        datalist (Sequence[xr.DataArray]): Data series from the cross-correlation experiments.
+        delayline_dim: the dimension name for "delay line", which must be in key of data.attrs
+            When this is the "position" dimention, the unit is assumed to be "mm". If the value has
+            already been converted to "time" dimension, set convert_position_to_time=True
+        delayline_origin (float): The value corresponding to the delay zero.
+        convert_position_to_time: (bool) If true, no conversion into "delay" is processed.
+
+    Returns: xr.DataArray
+    """
+    cross_correlations = []
+
+    for spectrum in datalist:
+        spectrum_arr = (
+            spectrum if isinstance(spectrum, xr.DataArray) else normalize_to_spectrum(spectrum)
+        )
+        if convert_position_to_time:
+            delay_time = spectrum_arr.attrs[delayline_dim] - delayline_origin
+        else:
+            delay_time = position_to_delaytime(
+                float(spectrum_arr[delayline_dim]),
+                delayline_origin,
+            )
+        cross_correlations.append(
+            spectrum_arr.assign_coords({"delay": delay_time}).expand_dims("delay"),
+        )
+    return xr.concat(cross_correlations, dim="delay")
+
+
+def delaytime_fs(mirror_movement_um: A) -> A:
+    """Return delaytime from the mirror movement (not position).
+
+    Args:
+        mirror_movement_um (float): mirror movement in micron unit.
+
+    Returns: float
+        delay time in fs.
+
+    """
+    return 3.335640951981521 * mirror_movement_um
+
+
+def position_to_delaytime(position_mm: A, delayline_offset_mm: float) -> A:
+    """Return delay time from the mirror position.
+
+    Args:
+        position_mm (np.ndarray | float): mirror position
+        delayline_offset_mm (float): mirror position corresponding to the zero delay
+
+    Returns: np.ndarray | float
+        delay time in fs unit.
+
+    """
+    return delaytime_fs(2 * (position_mm - delayline_offset_mm) * 1000)
 
 
 @update_provenance("Normalized subtraction map")

src/arpes/constants.py

@@ -21,18 +21,19 @@
 METERS_PER_SECOND_PER_EV_ANGSTROM = (
     151927  # converts from eV * angstrom to meters/second velocity units
 )
-HBAR = 1.0545718176461565 * 10 ** (-34)
-HBAR_PER_EV = 6.582119569 * 10 ** (
-    -16
-)  # gives the energy lifetime relationship via tau = -hbar / np.imag(self_energy)
+HBAR = 1.0545718176461565e-34
+HBAR_PER_EV = 6.582119569509067e-16
+# gives the energy lifetime relationship via tau = -hbar / np.imag(self_energy)
+
+
 BARE_ELECTRON_MASS = 9.109383701e-31  # kg
 HBAR_SQ_EV_PER_ELECTRON_MASS = 0.475600805657  # hbar^2 / m0 in eV^2 s^2 / kg
 HBAR_SQ_EV_PER_ELECTRON_MASS_ANGSTROM_SQ = 7.619964  # (hbar^2) / (m0 * angstrom ^2) in eV
 
 K_BOLTZMANN_EV_KELVIN = 8.617333262145178e-5  # in units of eV / Kelvin
 K_BOLTZMANN_MEV_KELVIN = 1000 * K_BOLTZMANN_EV_KELVIN  # meV / Kelvin
 
-HC = 1239.84172  # in units of eV * nm
+HC = 1239.8419843320028  # in units of eV * nm
 
 HEX_ALPHABET = "ABCDEF0123456789"
 ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

src/arpes/laser.py

@@ -2,12 +2,45 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger
+from typing import TYPE_CHECKING, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .constants import HC
 
 if TYPE_CHECKING:
     import pint
 
-__all__ = ("electrons_per_pulse",)
+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
+
+
+__all__ = ("electrons_per_pulse", "wavelength_to_energy")
+
+A = TypeVar("A", NDArray[np.float64], float)
+
+
+def wavelength_to_energy(wavelength_nm: A) -> A:
+    """Return Energy of the light.
+
+    Args:
+        wavelength_nm (NDArray | float): wavelength of the light in nm unit.
+
+    Returns: NDArray | float
+        Photon energy in eV unit.
+    """
+    return HC / wavelength_nm
 
 
 def electrons_per_pulse(

src/arpes/plotting/bands.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Unpack
 
 import matplotlib.pyplot as plt
+from matplotlib import colorbar
 from matplotlib.axes import Axes
 
 from arpes.provenance import save_plot_provenance
@@ -42,7 +43,6 @@ def plot_with_bands(
         out: [TODO:description]
         kwargs: pass to data.plot()
 
-
     Returns:
         [TODO:description]
     """
@@ -54,7 +54,9 @@ def plot_with_bands(
         title = data.S.label.replace("_", " ")
 
     mesh: AxesImage = data.plot(ax=ax, **kwargs)
-    mesh.colorbar.set_label(label_for_colorbar(data))
+    mesh_colorbar = mesh.colorbar
+    assert isinstance(mesh_colorbar, colorbar.Colorbar)
+    mesh_colorbar.set_label(label_for_colorbar(data))
 
     if data.S.is_differentiated:
         mesh.set_cmap("Blues")

src/arpes/plotting/bz.py

@@ -534,9 +534,8 @@ def draw(self, renderer) -> None:
 def annotate_special_paths(
     ax: Axes,
     paths: list[str] | str,
-    cell: Sequence[Sequence[float]] | None = None,
-    transformations=None,
-    offset: dict[str, float | Sequence[float]] | None = None,
+    cell: NDArray[np.float_] | Sequence[Sequence[float]] | None = None,
+    offset: dict[str, Sequence[float]] | None = None,
     special_points: dict[str, NDArray[np.float_]] | None = None,
     labels=None,
     **kwargs: Incomplete,
@@ -560,7 +559,7 @@ def annotate_special_paths(
         if labels is None:
             labels = paths
 
-        converted_paths = process_kpath(paths, cell, special_points=special_points)
+        converted_paths = process_kpath(paths, np.array(cell), special_points=special_points)
         logger.debug(f"converted_paths: {converted_paths}")
 
         if not isinstance(labels[0], list):