Skip to content

Commit

Permalink
💬 Update type hints
Browse files Browse the repository at this point in the history
  • Loading branch information
arafune committed Feb 1, 2024
1 parent 70fa5b0 commit e3220e8
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 13 deletions.
14 changes: 9 additions & 5 deletions arpes/analysis/band_analysis.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Provides some band analysis tools."""

from __future__ import annotations

import contextlib
Expand All @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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.
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down
13 changes: 5 additions & 8 deletions arpes/analysis/deconvolution.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Provides deconvolution implementations, especially for 2D Richardson-Lucy."""

from __future__ import annotations

import contextlib
Expand Down Expand Up @@ -46,19 +47,15 @@ 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
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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down
1 change: 1 addition & 0 deletions arpes/analysis/mask.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
"""Utilities for applying masks to data."""

from __future__ import annotations

from typing import TYPE_CHECKING
Expand Down

0 comments on commit e3220e8

Please sign in to comment.