diff --git a/gridlock/draw.py b/gridlock/draw.py index b2a3245..2322a66 100644 --- a/gridlock/draw.py +++ b/gridlock/draw.py @@ -1,7 +1,7 @@ """ Drawing-related methods for Grid class """ -from typing import List, Optional, Union, Sequence, Callable +from typing import Union, Sequence, Callable import numpy from numpy.typing import NDArray, ArrayLike @@ -27,7 +27,7 @@ def draw_polygons( center: ArrayLike, polygons: Sequence[NDArray], thickness: float, - foreground: Union[Sequence[foreground_t], foreground_t], + foreground: Sequence[foreground_t] | foreground_t, ) -> None: """ Draw polygons on an axis-aligned plane. @@ -70,7 +70,7 @@ def draw_polygons( + 'xyz'[surface_normal]) # Broadcast foreground where necessary - foregrounds: Union[Sequence[foreground_callable_t], Sequence[float]] + foregrounds: Sequence[foreground_callable_t] | Sequence[float] if numpy.size(foreground) == 1: # type: ignore foregrounds = [foreground] * len(cell_data) # type: ignore elif isinstance(foreground, numpy.ndarray): @@ -202,7 +202,7 @@ def draw_polygon( center: ArrayLike, polygon: ArrayLike, thickness: float, - foreground: Union[Sequence[foreground_t], foreground_t], + foreground: Sequence[foreground_t] | foreground_t, ) -> None: """ Draw a polygon on an axis-aligned plane. @@ -226,7 +226,7 @@ def draw_slab( surface_normal: int, center: ArrayLike, thickness: float, - foreground: Union[Sequence[foreground_t], foreground_t], + foreground: Sequence[foreground_t] | foreground_t, ) -> None: """ Draw an axis-aligned infinite slab. @@ -276,7 +276,7 @@ def draw_cuboid( cell_data: NDArray, center: ArrayLike, dimensions: ArrayLike, - foreground: Union[Sequence[foreground_t], foreground_t], + foreground: Sequence[foreground_t] | foreground_t, ) -> None: """ Draw an axis-aligned cuboid @@ -305,7 +305,7 @@ def draw_cylinder( radius: float, thickness: float, num_points: int, - foreground: Union[Sequence[foreground_t], foreground_t], + foreground: Sequence[foreground_t] | foreground_t, ) -> None: """ Draw an axis-aligned cylinder. Approximated by a num_points-gon diff --git a/gridlock/grid.py b/gridlock/grid.py index 5f132a0..36e0608 100644 --- a/gridlock/grid.py +++ b/gridlock/grid.py @@ -1,4 +1,4 @@ -from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar, TypeVar +from typing import Callable, Sequence, ClassVar, Self import numpy from numpy.typing import NDArray, ArrayLike @@ -12,7 +12,6 @@ from . import GridError foreground_callable_type = Callable[[NDArray, NDArray, NDArray], NDArray] -T = TypeVar('T', bound='Grid') class Grid: @@ -49,10 +48,10 @@ class Grid: Because of this, we either assume this 'ghost' cell is the same size as the last real cell, or, if `self.periodic[a]` is set to `True`, the same size as the first cell. """ - exyz: List[NDArray] + exyz: list[NDArray] """Cell edges. Monotonically increasing without duplicates.""" - periodic: List[bool] + periodic: list[bool] """For each axis, determines how far the rightmost boundary gets shifted. """ shifts: NDArray @@ -80,7 +79,7 @@ class Grid: from .position import ind2pos, pos2ind @property - def dxyz(self) -> List[NDArray]: + def dxyz(self) -> list[NDArray]: """ Cell sizes for each axis, no shifts applied @@ -90,7 +89,7 @@ class Grid: return [numpy.diff(ee) for ee in self.exyz] @property - def xyz(self) -> List[NDArray]: + def xyz(self) -> list[NDArray]: """ Cell centers for each axis, no shifts applied @@ -124,7 +123,7 @@ class Grid: return numpy.hstack((self.num_grids, self.shape)) @property - def dxyz_with_ghost(self) -> List[NDArray]: + def dxyz_with_ghost(self) -> list[NDArray]: """ Gives dxyz with an additional 'ghost' cell at the end, whose value depends on whether or not the axis has periodic boundary conditions. See main description @@ -153,7 +152,7 @@ class Grid: return numpy.array(centers, dtype=float) @property - def dxyz_limits(self) -> Tuple[NDArray, NDArray]: + def dxyz_limits(self) -> tuple[NDArray, NDArray]: """ Returns the minimum and maximum cell size for each axis, as a tuple of two 3-element ndarrays. No shifts are applied, so these are extreme bounds on these values (as a @@ -166,7 +165,7 @@ class Grid: d_max = numpy.array([max(self.dxyz[a]) for a in range(3)], dtype=float) return d_min, d_max - def shifted_exyz(self, which_shifts: Optional[int]) -> List[NDArray]: + def shifted_exyz(self, which_shifts: int | None) -> list[NDArray]: """ Returns edges for which_shifts. @@ -188,7 +187,7 @@ class Grid: return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)] - def shifted_dxyz(self, which_shifts: Optional[int]) -> List[NDArray]: + def shifted_dxyz(self, which_shifts: int | None) -> list[NDArray]: """ Returns cell sizes for `which_shifts`. @@ -215,7 +214,7 @@ class Grid: return sdxyz - def shifted_xyz(self, which_shifts: Optional[int]) -> List[NDArray[numpy.float64]]: + def shifted_xyz(self, which_shifts: int | None) -> list[NDArray[numpy.float64]]: """ Returns cell centers for `which_shifts`. @@ -231,7 +230,7 @@ class Grid: dxyz = self.shifted_dxyz(which_shifts) return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)] - def autoshifted_dxyz(self) -> List[NDArray[numpy.float64]]: + def autoshifted_dxyz(self) -> list[NDArray[numpy.float64]]: """ Return cell widths, with each dimension shifted by the corresponding shifts. @@ -242,7 +241,7 @@ class Grid: raise GridError('Autoshifting requires exactly 3 grids') return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)] - def allocate(self, fill_value: Optional[float] = 1.0, dtype=numpy.float32) -> NDArray: + def allocate(self, fill_value: float | None = 1.0, dtype=numpy.float32) -> NDArray: """ Allocate an ndarray for storing grid data. @@ -263,7 +262,7 @@ class Grid: self, pixel_edge_coordinates: Sequence[ArrayLike], shifts: ArrayLike = Yee_Shifts_E, - periodic: Union[bool, Sequence[bool]] = False, + periodic: bool | Sequence[bool] = False, ) -> None: """ Args: @@ -320,7 +319,7 @@ class Grid: g.__dict__.update(tmp_dict) return g - def save(self: T, filename: str) -> T: + def save(self, filename: str) -> Self: """ Save to file. @@ -334,7 +333,7 @@ class Grid: pickle.dump(self.__dict__, f, protocol=2) return self - def copy(self: T) -> T: + def copy(self) -> Self: """ Returns: Deep copy of the grid. diff --git a/gridlock/position.py b/gridlock/position.py index 35ec1df..b674716 100644 --- a/gridlock/position.py +++ b/gridlock/position.py @@ -1,8 +1,6 @@ """ Position-related methods for Grid class """ -from typing import List, Optional, Sequence - import numpy from numpy.typing import NDArray, ArrayLike @@ -12,7 +10,7 @@ from . import GridError def ind2pos( self, ind: NDArray, - which_shifts: Optional[int] = None, + which_shifts: int | None = None, round_ind: bool = True, check_bounds: bool = True ) -> NDArray[numpy.float64]: @@ -64,7 +62,7 @@ def ind2pos( def pos2ind( self, r: ArrayLike, - which_shifts: Optional[int], + which_shifts: int | None, round_ind: bool = True, check_bounds: bool = True ) -> NDArray[numpy.float64]: diff --git a/gridlock/read.py b/gridlock/read.py index 055670c..8ea4b50 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -1,13 +1,18 @@ """ Readback and visualization methods for Grid class """ -from typing import Dict, Optional, Union, Any +from typing import Any, TYPE_CHECKING import numpy from numpy.typing import NDArray, ArrayLike from . import GridError +if TYPE_CHECKING: + import matplotlib.axes + import matplotlib.figure + + # .visualize_* uses matplotlib # .visualize_isosurface uses skimage # .visualize_isosurface uses mpl_toolkits.mplot3d @@ -128,7 +133,7 @@ def visualize_slice( def visualize_isosurface( self, cell_data: NDArray, - level: Optional[float] = None, + level: float | None = None, which_shifts: int = 0, sample_period: int = 1, show_edges: bool = True,