Skip to content

Commit

Permalink
🐛 fix bug in xarray_extensions occuring when temperature is not numer…
Browse files Browse the repository at this point in the history
…ic, such as "RT", "LT".

💬  update type hints
  • Loading branch information
arafune committed Oct 19, 2023
1 parent ec0105f commit fe42f5c
Show file tree
Hide file tree
Showing 7 changed files with 74 additions and 21 deletions.
9 changes: 7 additions & 2 deletions arpes/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,8 +217,13 @@ class SCANINFO(TypedDict, total=False):
sample: str | None


class EXPERIMENTALINFO(TypedDict, total=False):
temperature: float | None
class ExperimentalConditions(TypedDict, total=True):
hv: float
polarization: float | str
temperature: float | str


class EXPERIMENTALINFO(ExperimentalConditions, total=False):
temperature_cryotip: float | None
pressure: float | None
polarization: float | tuple[float | None, float | None] | str
Expand Down
5 changes: 3 additions & 2 deletions arpes/endstations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@
]

LOGLEVELS = (DEBUG, INFO)
LOGLEVEL = LOGLEVELS[0]
LOGLEVEL = LOGLEVELS[1]
logger = getLogger(__name__)
fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s"
formatter = Formatter(fmt)
Expand Down Expand Up @@ -1061,11 +1061,12 @@ def resolve_endstation(*, retry: bool = True, **kwargs: Incomplete) -> type:
if endstation_name is None:
warnings.warn("Endstation not provided. Using `fallback` plugin.", stacklevel=2)
endstation_name = "fallback"

logger.debug(f"_ENDSTATION_ALIASES is : {_ENDSTATION_ALIASES}")
try:
return endstation_from_alias(endstation_name)
except KeyError:
if retry:
logger.debug("KeyError occurred")
import arpes.config # pylint: disable=redefined-outer-name

arpes.config.load_plugins()
Expand Down
2 changes: 1 addition & 1 deletion arpes/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@


LOGLEVELS = (DEBUG, INFO)
LOGLEVEL = LOGLEVELS[0]
LOGLEVEL = LOGLEVELS[1]
logger = getLogger(__name__)
fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s"
formatter = Formatter(fmt)
Expand Down
46 changes: 37 additions & 9 deletions arpes/plotting/annotations.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
"""Annotations onto plots for experimental conditions or locations."""
from __future__ import annotations

from typing import TYPE_CHECKING, Unpack
from typing import TYPE_CHECKING, Literal, Unpack

import numpy as np
import xarray as xr
import matplotlib as mpl
from matplotlib.axes import Axes
from mpl_toolkits.mplot3d import Axes3D

Expand All @@ -18,7 +19,7 @@
from _typeshed import Incomplete
from numpy.typing import NDArray

from arpes._typing import DataType, MPLTextParam
from arpes._typing import DataType, ExperimentalConditions, MPLTextParam

__all__ = (
"annotate_cuts",
Expand All @@ -37,7 +38,7 @@ def annotate_experimental_conditions(
desc: list[str | float] | float | str,
*,
show: bool = False,
orientation: str = "top",
orientation: Literal["top", "bottom"] = "top",
**kwargs: Unpack[MPLTextParam],
) -> None:
"""Renders information about the experimental conditions onto a set of axes.
Expand Down Expand Up @@ -65,20 +66,47 @@ def annotate_experimental_conditions(
ax.set_axis_off()
ax.patch.set_alpha(0)

delta = -1
delta: float = -1
current = 100.0
if orientation == "bottom":
delta = 1
current = 0

fontsize = kwargs.get("fontsize", 16)
fontsize_keyword = kwargs.get("fontsize", 16)
if isinstance(fontsize_keyword, float):
fontsize = fontsize_keyword
elif fontsize_keyword in (
"xx-small",
"x-small",
"small",
"medium",
"large",
"x-large",
"xx-large",
):
font_scalings = { # see matplotlib.font_manager
"xx-small": 0.579,
"x-small": 0.694,
"small": 0.833,
"medium": 1.0,
"large": 1.200,
"x-large": 1.440,
"xx-large": 1.728,
"larger": 1.2,
"smaller": 0.833,
}
fontsize = mpl.rc_params["font.size"] * font_scalings[fontsize_keyword]
else:
err_msg = "Incorrect font size setting"
raise RuntimeError(err_msg)
delta = fontsize * delta

conditions = data.S.experimental_conditions
conditions: ExperimentalConditions = data.S.experimental_conditions

renderers = {
"temp": lambda c: "\\textbf{T = " + "{:.3g}".format(c["temp"]) + " K}",
"photon": _render_photon,
"hv": _render_photon,
"photon polarization": lambda c: _render_photon(c) + ", " + _render_polarization(c),
"polarization": _render_polarization,
}
Expand All @@ -88,9 +116,9 @@ def annotate_experimental_conditions(
current += item + delta
continue

item = item.replace("_", " ").lower()
item_replaced = item.replace("_", " ").lower()

ax.text(0, current, renderers[item](conditions), **kwargs)
ax.text(0, current, renderers[item_replaced](conditions), **kwargs)
current += delta


Expand Down Expand Up @@ -118,7 +146,7 @@ def _render_polarization(conditions: dict[str, str]) -> str:


def _render_photon(c: dict[str, float]) -> str:
return "\\textbf{" + str(c["hv"]) + " eV"
return "\\textbf{" + str(c["hv"]) + " eV}"


def annotate_cuts(
Expand Down
13 changes: 11 additions & 2 deletions arpes/plotting/bz.py
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,11 @@ def plot_plane_to_bz(
ax.add_collection3d(collection, zs="z")


def plot_data_to_bz(data: DataType, cell: Sequence[Sequence[float]], **kwargs: Incomplete):
def plot_data_to_bz(
data: DataType,
cell: Sequence[Sequence[float]],
**kwargs: Incomplete,
):
"""A dimension agnostic tool used to plot ARPES data onto a Brillouin zone."""
if len(data) == 3: # noqa: PLR2004
return plot_data_to_bz3d(data, cell, **kwargs)
Expand Down Expand Up @@ -307,6 +311,11 @@ def plot_data_to_bz3d(
) -> None:
"""Plots ARPES data onto a 3D Brillouin zone."""
msg = "plot_data_to_bz3d is not implemented yet."
logger.debug(f"id of data: {data.attrs.get('id', None)}")
logger.debug(f"cell: {cell}")
if kwargs:
for k, v in kwargs.items():
logger.debug(f"kwargs; k: {k}, v: {v}")
raise NotImplementedError(msg)


Expand Down Expand Up @@ -674,7 +683,7 @@ def bz2d_plot(
logger.debug(f"paths: {paths}")
logger.debug(f"points: {points}")
logger.debug(f"repeat: {repeat}")
logger.debug(f"transfomations: {transformations}")
logger.debug(f"transformations: {transformations}")
logger.debug(f"hide_ax: {hide_ax}")
logger.debug(f"vectors: {vectors}")
logger.debug(f"set_equal_aspect: {set_equal_aspect}")
Expand Down
18 changes: 14 additions & 4 deletions arpes/xarray_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,13 @@
from matplotlib.typing import RGBColorType
from numpy.typing import DTypeLike, NDArray

from arpes._typing import ANGLE, SPECTROMETER, DataType, PColorMeshKwargs
from arpes._typing import (
ANGLE,
SPECTROMETER,
DataType,
ExperimentalConditions,
PColorMeshKwargs,
)

IncompleteMPL: TypeAlias = Incomplete

Expand All @@ -103,7 +109,7 @@
ANGLE_VARS = ("alpha", "beta", "chi", "psi", "phi", "theta")

LOGLEVELS = (DEBUG, INFO)
LOGLEVEL = LOGLEVELS[0]
LOGLEVEL = LOGLEVELS[1]
logger = getLogger(__name__)
fmt = "%(asctime)s %(levelname)s %(name)s :%(message)s"
formatter = Formatter(fmt)
Expand Down Expand Up @@ -160,7 +166,7 @@ def sherman_function(self) -> Incomplete:
raise ValueError(msg)

@property
def experimental_conditions(self) -> dict[str, str | float | None]:
def experimental_conditions(self) -> ExperimentalConditions:
"""Return experimental condition: hv, polarization, temperature.
Use this property in plotting/annotations.py/conditions
Expand Down Expand Up @@ -190,6 +196,8 @@ def polarization(self) -> str | None:
}.get(int(self._obj.attrs["epu_pol"]))
except ValueError:
return self._obj.attrs["epu_pol"]
if "pol" in self._obj.attrs:
return self._obj.attrs["pol"]
return None

@property
Expand Down Expand Up @@ -1820,6 +1828,7 @@ def _repr_html_spectrometer_info(self) -> str:

@staticmethod
def _repr_html_experimental_conditions(conditions: dict[str, str | float | None]) -> str:
logger.debug(f"conditions: {conditions}")
transforms = {
"polarization": lambda p: {
"p": "Linear Horizontal",
Expand All @@ -1842,9 +1851,10 @@ def no_change(x: str | float) -> str | float:
for k, v in conditions.items():
if v is None:
continue
if np.isnan(v):
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:
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ omit = [
exclude_lines = ["raise NotImplementedError", "pragma: no cover"]

[tool.pyright]
typeCheckingMode = "basic"
typeCheckingMode = "off"

pythonVersion = "3.11"
pythonPlatform = "All"
Expand Down

0 comments on commit fe42f5c

Please sign in to comment.