From a1c6585c5c02463c6cacdd43d37fc181050a4b9f Mon Sep 17 00:00:00 2001 From: Ryuichi Arafune Date: Thu, 1 Feb 2024 20:42:25 +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/_typing.py | 2 +- arpes/xarray_extensions.py | 42 ++++++++++++++++----------------- tests/test_xarray_extensions.py | 2 +- 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/arpes/_typing.py b/arpes/_typing.py index c9fa0c8d..4c12e7f6 100644 --- a/arpes/_typing.py +++ b/arpes/_typing.py @@ -221,7 +221,7 @@ class _BEAMLINEINFO(TypedDict, total=False): class LIGHTSOURCEINFO(_PROBEINFO, _PUMPINFO, _BEAMLINEINFO, total=False): - polarization: float | tuple[float | None, float | None] | str + polarization: float | tuple[float, float] | str photon_flux: float photocurrent: float probe: None diff --git a/arpes/xarray_extensions.py b/arpes/xarray_extensions.py index adefe90c..86a47369 100644 --- a/arpes/xarray_extensions.py +++ b/arpes/xarray_extensions.py @@ -203,7 +203,9 @@ def sherman_function(self) -> Incomplete: raise ValueError(msg) @property - def experimental_conditions(self) -> dict[str, float | str]: + def experimental_conditions( + self, + ) -> EXPERIMENTINFO: """Return experimental condition: hv, polarization, temperature. Use this property in plotting/annotations.py/conditions @@ -215,7 +217,7 @@ def experimental_conditions(self) -> dict[str, float | str]: } @property - def polarization(self) -> float | str | None: + def polarization(self) -> float | str | tuple[float, float]: """Returns the light polarization information. ToDo: Test @@ -228,12 +230,12 @@ def polarization(self) -> float | str | None: 0: "p", 1: "rc", 2: "s", - }.get(int(self._obj.attrs["epu_pol"])) + }.get(int(self._obj.attrs["epu_pol"]), np.nan) except ValueError: return self._obj.attrs["epu_pol"] if "pol" in self._obj.attrs: return self._obj.attrs["pol"] - return None + return np.nan @property def is_subtracted(self) -> bool: @@ -951,7 +953,9 @@ def find_spectrum_energy_edges(self, *, indices: bool = False) -> NDArray[np.flo energy_marginal = self._obj.sum([d for d in self._obj.dims if d not in ["eV"]]) embed_size = 20 - embedded = np.ndarray(shape=[embed_size, *list(energy_marginal.values.shape)]) + embedded: NDArray[np.float_] = np.ndarray( + shape=[embed_size, *list(energy_marginal.values.shape)], + ) embedded[:] = energy_marginal.values embedded = ndi.gaussian_filter(embedded, embed_size / 3) @@ -971,17 +975,17 @@ def find_spectrum_angular_edges_full( self, *, indices: bool = False, + energy_division: float = 0.05, ) -> tuple[NDArray[np.float_], NDArray[np.float_], xr.DataArray]: # as a first pass, we need to find the bottom of the spectrum, we will use this # to select the active region and then to rebin into course steps in energy from 0 # down to this region # we will then find the appropriate edge for each slice, and do a fit to the edge locations + energy_edge: NDArray[np.float_] = self.find_spectrum_energy_edges() + low_edge = np.min(energy_edge) + energy_division + high_edge = np.max(energy_edge) - energy_division - energy_edge = self.find_spectrum_energy_edges() - low_edge = np.min(energy_edge) + 0.05 - high_edge = np.max(energy_edge) - 0.05 - - if high_edge - low_edge < 0.15: # noqa: PLR2004 + if high_edge - low_edge < 3 * energy_division: # Doesn't look like the automatic inference of the energy edge was valid high_edge = np.max(self._obj.coords["eV"].values) low_edge = np.min(self._obj.coords["eV"].values) @@ -989,7 +993,7 @@ def find_spectrum_angular_edges_full( angular_dim = "pixel" if "pixel" in self._obj.dims else "phi" energy_cut = self._obj.sel(eV=slice(low_edge, high_edge)).S.sum_other(["eV", angular_dim]) - n_cuts = int(np.ceil((high_edge - low_edge) / 0.05)) + n_cuts = int(np.ceil((high_edge - low_edge) / energy_division)) new_shape = {"eV": n_cuts} new_shape[angular_dim] = len(energy_cut.coords[angular_dim].values) rebinned = rebin(energy_cut, shape=new_shape) @@ -1126,7 +1130,7 @@ def find_spectrum_angular_edges( ) embed_size = 20 - embedded = np.ndarray(shape=[embed_size, *list(near_ef.values.shape)]) + embedded: NDArray[np.float_] = np.ndarray(shape=[embed_size, *list(near_ef.values.shape)]) embedded[:] = near_ef.values embedded = ndi.gaussian_filter(embedded, embed_size / 3) @@ -1751,18 +1755,15 @@ def _repr_html_spectrometer_info(self) -> str: return ARPESAccessorBase.dict_to_html(ordered_settings) @staticmethod - def _repr_html_experimental_conditions(conditions: dict[str, str | float | None]) -> str: - """[TODO:summary]. + def _repr_html_experimental_conditions(conditions: EXPERIMENTINFO) -> str: + """Return the experimental conditions with html format. Args: - conditions: [TODO:description] - - Returns: - [TODO:description] + conditions (EXPERIMENTINFO): self.confitions is usually used. - Todo: Remove pandas related routine. (xarray including pasdas object is not so common.) + Returns (str): + html representation of the experimental conditions. """ - logger.debug(f"conditions: {conditions}") transforms = { "polarization": lambda p: { "p": "Linear Horizontal", @@ -1788,7 +1789,6 @@ def no_change(x: str | float) -> str | float: if isinstance(v, float) and np.isnan(v): continue transformed_dict[str(k)] = transforms.get(k, no_change)(v) - logger.debug(f"transformed_dict: {transformed_dict}") return ARPESAccessorBase.dict_to_html(transformed_dict) def _repr_html_(self) -> str: diff --git a/tests/test_xarray_extensions.py b/tests/test_xarray_extensions.py index 7d45e7aa..9cd942e3 100644 --- a/tests/test_xarray_extensions.py +++ b/tests/test_xarray_extensions.py @@ -146,7 +146,7 @@ def test_experimental_conditions(self, dataset_cut: xr.Dataset) -> None: """Test for property experimenta_conditions.""" assert dataset_cut.S.experimental_conditions == { "hv": 5.93, - "polarization": None, + "polarization": np.nan, "temperature": np.nan, }