diff --git a/Changes.md b/Changes.md index 1ac1576b..392e8092 100644 --- a/Changes.md +++ b/Changes.md @@ -2,7 +2,7 @@ Major Changes from 3.0.1 - Most important change: - - Use correct method to convert from the angle to moementum. (Original approach was found to be incorrect) + - Use correct method to convert from the angle to momentum. (The original way was incorrect) - New feature - Provide SPD_main.py & prodigy_itx.py diff --git a/README.rst b/README.rst index 969122c7..9c3a9375 100644 --- a/README.rst +++ b/README.rst @@ -18,8 +18,8 @@ :target: https://github.com/arafune/arpes/actions/workflows/test.yml .. |code style| image:: https://img.shields.io/badge/code%20style-black-000000.svg :target: https://github.com/psf/black - - +.. |code fromat| image:: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json + :target: https://github.com/astral-sh/ruff PyARPES ======= diff --git a/src/arpes/analysis/deconvolution.py b/src/arpes/analysis/deconvolution.py index 97db4d04..5a94658b 100644 --- a/src/arpes/analysis/deconvolution.py +++ b/src/arpes/analysis/deconvolution.py @@ -2,12 +2,14 @@ from __future__ import annotations +from logging import DEBUG, INFO, Formatter, StreamHandler, getLogger from typing import TYPE_CHECKING import numpy as np import scipy import scipy.ndimage import xarray as xr +from scipy.stats import multivariate_normal from skimage.restoration import richardson_lucy import arpes.xarray_extensions # noqa: F401 @@ -16,6 +18,8 @@ from arpes.utilities import normalize_to_spectrum if TYPE_CHECKING: + from collections.abc import Hashable + from numpy.typing import NDArray @@ -25,7 +29,17 @@ "make_psf1d", ) -TWO_DIMWENSION = 2 +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("Approximate Iterative Deconvolution") @@ -120,15 +134,61 @@ def make_psf1d(data: xr.DataArray, dim: str, sigma: float) -> xr.DataArray: @update_provenance("Make Point Spread Function") -def make_psf(data: xr.DataArray, sigmas: dict[str, float]) -> xr.DataArray: +def make_psf( + data: xr.DataArray, + sigmas: dict[Hashable, float], + *, + fwhm: bool = True, +) -> xr.DataArray: """Produces an n-dimensional gaussian point spread function for use in deconvolve_rl. - Not yet operational. - Args: data (DataType): input data sigmas (dict[str, float]): sigma values for each dimension. + fwhm (bool): if True, sigma is FWHM, not the standard deviation. Returns: The PSF to use. """ + strides = data.G.stride(generic_dim_names=False) + logger.debug(f"strides: {strides}") + assert set(strides) == set(sigmas) + pixels: dict[Hashable, int] = dict( + zip( + data.dims, + tuple([i - 1 if i % 2 == 0 else i for i in data.shape]), + strict=True, + ), + ) + + if fwhm: + sigmas = {k: v / (2 * np.sqrt(2 * np.log(2))) for k, v, in sigmas.items()} + cov: NDArray[np.float_] = np.zeros((len(sigmas), len(sigmas))) + for i, dim in enumerate(data.dims): + cov[i][i] = sigmas[dim] ** 2 # sigma is deviation, but multivariate_normal uses covariant + logger.debug(f"cov: {cov}") + + psf_coords: dict[Hashable, NDArray[np.float_]] = {} + for k in data.dims: + psf_coords[str(k)] = np.linspace( + -(pixels[str(k)] - 1) / 2 * strides[str(k)], + (pixels[str(k)] - 1) / 2 * strides[str(k)], + pixels[str(k)], + ) + if LOGLEVEL == DEBUG: + for k, v in psf_coords.items(): + logger.debug( + 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}") + return xr.DataArray( + multivariate_normal(mean=np.zeros(len(sigmas)), cov=cov).pdf( + coords_for_pdf_pos, + ), + dims=data.dims, + coords=psf_coords, + name="PSF", + ) diff --git a/src/arpes/endstations/__init__.py b/src/arpes/endstations/__init__.py index fe1d3746..172d6b2b 100644 --- a/src/arpes/endstations/__init__.py +++ b/src/arpes/endstations/__init__.py @@ -97,7 +97,7 @@ class EndstationBase: ALIASES: ClassVar[list[str]] = [] PRINCIPAL_NAME = "" - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = {} + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = {} MERGE_ATTRS: ClassVar[SPECTROMETER] = {} _SEARCH_DIRECTORIES: tuple[str, ...] = ( @@ -176,7 +176,7 @@ def is_file_accepted( return False try: - _ = cls.find_first_file(str(file)) + _ = cls.find_first_file(int(file)) except ValueError: return False return True diff --git a/src/arpes/endstations/plugin/ALG_main.py b/src/arpes/endstations/plugin/ALG_main.py index d64d2f5b..4607fd16 100644 --- a/src/arpes/endstations/plugin/ALG_main.py +++ b/src/arpes/endstations/plugin/ALG_main.py @@ -37,7 +37,7 @@ class ALGMainChamber(HemisphericalEndstation, FITSEndstation): "ALG-Main Chamber", ] - ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, list[str] | str]]]] = { + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = { "START_T": lambda _: {"time": " ".join(_.split(" ")[1:]).lower(), "date": _.split(" ")[0]}, } diff --git a/src/arpes/endstations/plugin/BL10_SARPES.py b/src/arpes/endstations/plugin/BL10_SARPES.py index 965b5cc0..bc143285 100644 --- a/src/arpes/endstations/plugin/BL10_SARPES.py +++ b/src/arpes/endstations/plugin/BL10_SARPES.py @@ -55,7 +55,7 @@ class BL10012SARPESEndstation(SynchrotronEndstation, HemisphericalEndstation, SE # Look at merlin.py for details } - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = { + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = { # TODO: Kayla or another user should add these # Look at merlin.py for details } diff --git a/src/arpes/endstations/plugin/MAESTRO.py b/src/arpes/endstations/plugin/MAESTRO.py index 453936d3..dc779f7a 100644 --- a/src/arpes/endstations/plugin/MAESTRO.py +++ b/src/arpes/endstations/plugin/MAESTRO.py @@ -138,7 +138,7 @@ class MAESTROMicroARPESEndstation(MAESTROARPESEndstationBase): "Z": "z", } - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = { + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, int | list[str] | str]]]] = { "START_T": lambda _: { "time": " ".join(_.split(" ")[1:]).lower(), "date": _.split(" ")[0], @@ -255,7 +255,7 @@ class MAESTRONanoARPESEndstation(MAESTROARPESEndstationBase): "Slit Defl.": "psi", } - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = { + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = { "START_T": lambda _: { "time": " ".join(_.split(" ")[1:]).lower(), "date": _.split(" ")[0], diff --git a/src/arpes/endstations/plugin/example_data.py b/src/arpes/endstations/plugin/example_data.py index f27a5a47..41f5b6cc 100644 --- a/src/arpes/endstations/plugin/example_data.py +++ b/src/arpes/endstations/plugin/example_data.py @@ -71,7 +71,7 @@ def load_single_frame( if len(coord.values.shape) and cname not in data.dims: replacement_coords[cname] = coord.mean().item() - data = data.assign_coords(**replacement_coords) + data = data.assign_coords(replacement_coords) # Wrap into a dataset dataset = xr.Dataset({"spectrum": data}) diff --git a/src/arpes/endstations/plugin/igor_plugin.py b/src/arpes/endstations/plugin/igor_plugin.py index 824e7cc8..2ee2279e 100644 --- a/src/arpes/endstations/plugin/igor_plugin.py +++ b/src/arpes/endstations/plugin/igor_plugin.py @@ -74,7 +74,7 @@ class IgorEndstation(SingleFileEndstation): MERGE_ATTRS: ClassVar[SPECTROMETER] = {} - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = {} + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = {} def load_single_frame( self, diff --git a/src/arpes/endstations/plugin/merlin.py b/src/arpes/endstations/plugin/merlin.py index 9aaca7ef..a6a0f336 100644 --- a/src/arpes/endstations/plugin/merlin.py +++ b/src/arpes/endstations/plugin/merlin.py @@ -109,7 +109,7 @@ class BL403ARPESEndstation(SynchrotronEndstation, HemisphericalEndstation, SESEn "undulator_type": "elliptically_polarized_undulator", } - ATTR_TRANSFORMS: ClassVar[dict[str, Callable]] = { + ATTR_TRANSFORMS: ClassVar[dict[str, Callable[..., dict[str, float | list[str] | str]]]] = { "acquisition_mode": lambda _: _.lower(), "lens_mode": lambda _: { "lens_mode": None, diff --git a/src/arpes/utilities/bz.py b/src/arpes/utilities/bz.py index a85c5613..0acccb0c 100644 --- a/src/arpes/utilities/bz.py +++ b/src/arpes/utilities/bz.py @@ -659,7 +659,7 @@ def reduced_bz_E_mask( selector = {} selector[data.dims[skip_col]] = selector_val - sdata = data.sel(**selector, method="nearest") + sdata = data.sel(selector, method="nearest") path = matplotlib.path.Path(poly_points) grid = np.array(