diff --git a/pyproject.toml b/pyproject.toml index f654a5e..89ed172 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,11 +8,12 @@ description = "SPH analysis and visualization" authors = [{name = "Terrence S. Tricco", email = "tstricco@mun.ca"}] readme = "README.md" dynamic = ["version"] -requires-python = ">=3.8" +requires-python = ">=3.10" dependencies = [ "matplotlib>=3.6", "numba>=0.55.1", - "pandas>=1.4", + "pandas>=2.0", + "numpy>=2.0", "scipy>=1.8.0", "seaborn>=0.11", ] diff --git a/sarracen/__init__.py b/sarracen/__init__.py index f30e21b..a47aafb 100644 --- a/sarracen/__init__.py +++ b/sarracen/__init__.py @@ -9,9 +9,9 @@ try: from .readers.read_shamrock_vtk import read_shamrock_vtk except ModuleNotFoundError as e: - # temporary workaround for no VTK in Python 3.8 on Mac ARM if e.name == "vtk": - def read_shamrock_vtk(*args, **kwargs): + def read_shamrock_vtk(filename: str, + pmass: float) -> 'SarracenDataFrame': raise ImportError("vtk is not installed.") else: raise diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 4070837..10914d2 100644 --- a/sarracen/disc/surface_density.py +++ b/sarracen/disc/surface_density.py @@ -4,21 +4,25 @@ from ..sarracen_dataframe import SarracenDataFrame from .utils import _get_mass, _get_origin from .utils import _bin_particles_by_radius, _get_bin_midpoints +from typing import TypeAlias -from typing import Tuple, Union + +ProfileResult: TypeAlias = np.ndarray | tuple[np.ndarray, np.ndarray] +SDResult: TypeAlias = (ProfileResult | + tuple[np.ndarray, np.ndarray, np.ndarray] | + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]) +AngularMomentumResult: TypeAlias = tuple[np.ndarray, ...] def azimuthal_average(data: 'SarracenDataFrame', target: str, - r_in: Union[float, None] = None, - r_out: Union[float, None] = None, + r_in: float | None = None, + r_out: float | None = None, bins: int = 300, log: bool = False, geometry: str = 'cylindrical', - origin: Union[list, None] = None, - retbins: bool = False) -> Union[np.ndarray, - Tuple[np.ndarray, - np.ndarray]]: + origin: list[float] | None = None, + retbins: bool = False) -> ProfileResult: """ Calculates the 1D azimuthally-averaged profile for a target quantity. @@ -44,7 +48,7 @@ def azimuthal_average(data: 'SarracenDataFrame', geometry : str, optional Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. retbins : bool, optional @@ -52,10 +56,10 @@ def azimuthal_average(data: 'SarracenDataFrame', Returns ------- - array + ndarray A NumPy array of length bins containing the averaged profile. - array, optional - The midpoint values of each bin. Only returned if *retbins=True*. + ndarray, optional + The midpoint values of each bin. Only returned if `retbins=True`. Raises ------ @@ -76,22 +80,13 @@ def azimuthal_average(data: 'SarracenDataFrame', def surface_density(data: 'SarracenDataFrame', - r_in: Union[float, None] = None, - r_out: Union[float, None] = None, + r_in: float | None = None, + r_out: float | None = None, bins: int = 300, log: bool = False, geometry: str = 'cylindrical', - origin: Union[list, None] = None, - retbins: bool = False) -> Union[np.ndarray, - Tuple[np.ndarray, - np.ndarray], - Tuple[np.ndarray, - np.ndarray, - np.ndarray], - Tuple[np.ndarray, - np.ndarray, - np.ndarray, - np.ndarray]]: + origin: list[float] | None = None, + retbins: bool = False) -> SDResult: """ Calculates the 1D azimuthally-averaged surface density profile. @@ -115,7 +110,7 @@ def surface_density(data: 'SarracenDataFrame', geometry : str, optional Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. retbins : bool, optional @@ -123,9 +118,14 @@ def surface_density(data: 'SarracenDataFrame', Returns ------- - array + ndarray or tuple of ndarray A NumPy array of length bins containing the surface density profile. - array, optional + For gas-only data, this is the gas surface density profile. For + one-fluid data, the gas and dust surface desity profiles are returned. + For multiple small dust grains, the gas surface density profile is + returned, along with the total dust profile and per-grain dust + profiles. + ndarray, optional The midpoint values of each bin. Only returned if *retbins=True*. Raises @@ -194,7 +194,7 @@ def surface_density(data: 'SarracenDataFrame', else: return sigma_gas, sigma_dust - elif ndustsmall > 1: # multiple grain sizes + else: # multiple grain sizes if 'dustfrac_total' not in data.columns: data.calc_one_fluid_quantities() @@ -221,28 +221,27 @@ def surface_density(data: 'SarracenDataFrame', def _calc_angular_momentum(data: 'SarracenDataFrame', rbins: pd.Series, - origin: list, - unit_vector: bool) -> Tuple[pd.Series, - pd.Series, + origin: list[float], + unit_vector: bool) -> tuple[pd.Series, pd.Series, pd.Series]: """ Utility function to calculate angular momentum of the disc. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame Particle data, in a SarracenDataFrame. - rbins: Series + rbins : Series The radial bin to which each particle belongs. - origin: list + origin : list of float The x, y and z centre point around which to compute radii. - unit_vector: bool + unit_vector : bool Whether to convert the angular momentum to unit vectors. Default is True. Returns ------- - Lx, Ly, Lz: Series + Lx, Ly, Lz : Series The x, y and z components of the angular momentum per bin. """ @@ -276,15 +275,14 @@ def _calc_angular_momentum(data: 'SarracenDataFrame', def angular_momentum(data: 'SarracenDataFrame', - r_in: Union[float, None] = None, - r_out: Union[float, None] = None, + r_in: float | None = None, + r_out: float | None = None, bins: int = 300, log: bool = False, geometry: str = 'cylindrical', - origin: Union[list, None] = None, + origin: list[float] | None = None, retbins: bool = False, - unit_vector: bool = True) -> Union[Tuple[np.ndarray, - ...]]: + unit_vector: bool = True) -> AngularMomentumResult: """ Calculates the angular momentum profile of the disc. @@ -308,21 +306,22 @@ def angular_momentum(data: 'SarracenDataFrame', geometry : str, optional Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. retbins : bool, optional Whether to return the midpoints of the bins or not. Defaults to False. - unit_vector: bool, optional + unit_vector : bool, optional Whether to convert the angular momentum to unit vectors. Default is True. Returns ------- - array - A NumPy array of length bins containing the angular momentum profile. - array, optional - The midpoint values of each bin. Only returned if *retbins=True*. + tuple of ndarray + Three NumPy arrays of length bins containing the Lx, Ly and Lz angular + momentum profiles. + ndarray, optional + The midpoint values of each bin. Only returned if `retbins=True`. Raises ------ @@ -350,17 +349,17 @@ def angular_momentum(data: 'SarracenDataFrame', def _calc_scale_height(data: 'SarracenDataFrame', rbins: pd.Series, - origin: list) -> pd.Series: + origin: list[float]) -> pd.Series: """ Utility function to calculate the scale height of the disc. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame Particle data, in a SarracenDataFrame. - rbins: Series + rbins : Series The radial bin to which each particle belongs. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. @@ -380,15 +379,13 @@ def _calc_scale_height(data: 'SarracenDataFrame', def scale_height(data: 'SarracenDataFrame', - r_in: Union[float, None] = None, - r_out: Union[float, None] = None, + r_in: float | None = None, + r_out: float | None = None, bins: int = 300, log: bool = False, geometry: str = 'cylindrical', - origin: Union[list, None] = None, - retbins: bool = False) -> Union[np.ndarray, - Tuple[np.ndarray, - np.ndarray]]: + origin: list[float] | None = None, + retbins: bool = False) -> ProfileResult: """ Calculates the scale height, H/R, of the disc. @@ -415,7 +412,7 @@ def scale_height(data: 'SarracenDataFrame', geometry : str, optional Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. retbins : bool, optional @@ -423,10 +420,10 @@ def scale_height(data: 'SarracenDataFrame', Returns ------- - array + ndarray A NumPy array of length bins scale height, H, profile. - array, optional - The midpoint values of each bin. Only returned if *retbins=True*. + ndarray, optional + The midpoint values of each bin. Only returned if `retbins=True`. Raises ------ @@ -452,14 +449,13 @@ def scale_height(data: 'SarracenDataFrame', def honH(data: 'SarracenDataFrame', - r_in: Union[float, None] = None, - r_out: Union[float, None] = None, + r_in: float | None = None, + r_out: float | None = None, bins: int = 300, log: bool = False, geometry: str = 'cylindrical', - origin: Union[list, None] = None, - retbins: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, - np.ndarray]]: + origin: list[float] | None = None, + retbins: bool = False) -> ProfileResult: """ Calculates /H, the averaged smoothing length divided by the scale height. @@ -484,7 +480,7 @@ def honH(data: 'SarracenDataFrame', geometry : str, optional Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like, optional + origin : array_like, optional The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. retbins : bool, optional @@ -492,10 +488,10 @@ def honH(data: 'SarracenDataFrame', Returns ------- - array + ndarray A NumPy array of length bins containing the /H profile. - array, optional - The midpoint values of each bin. Only returned if *retbins=True*. + ndarray, optional + The midpoint values of each bin. Only returned if `retbins=True`. Raises ------ diff --git a/sarracen/disc/utils.py b/sarracen/disc/utils.py index a6c2d99..ebec2c6 100644 --- a/sarracen/disc/utils.py +++ b/sarracen/disc/utils.py @@ -1,11 +1,11 @@ import numpy as np import pandas as pd import sys + from ..sarracen_dataframe import SarracenDataFrame -from typing import Tuple, Union -def _get_mass(data: 'SarracenDataFrame') -> Union[float, pd.Series]: +def _get_mass(data: 'SarracenDataFrame') -> float | pd.Series: if data.mcol is None: if 'mass' not in data.params: raise KeyError("'mass' column does not exist in this " @@ -15,26 +15,28 @@ def _get_mass(data: 'SarracenDataFrame') -> Union[float, pd.Series]: return data[data.mcol] -def _get_origin(origin: Union[list, None]) -> list: +def _get_origin(origin: list[float] | None) -> list[float]: if origin is None: return [0.0, 0.0, 0.0] else: - return origin + return list(origin) def _bin_particles_by_radius(data: 'SarracenDataFrame', - r_in: Union[float, None], - r_out: Union[float, None], + r_in: float | None, + r_out: float | None, bins: int, log: bool, geometry: str, - origin: list) -> Tuple[pd.Series, np.ndarray]: + origin: list[float]) -> tuple[pd.Series, + np.ndarray]: + """ Utility function to bin particles in discrete intervals by radius. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset. r_in : float Inner radius of the disc. Defaults to the minimum r value. @@ -48,15 +50,15 @@ def _bin_particles_by_radius(data: 'SarracenDataFrame', geometry : str Coordinate system to use to calculate the particle radii. Can be either *spherical* or *cylindrical*. Defaults to *cylindrical*. - origin : array-like + origin : array_like The x, y and z centre point around which to compute radii. Defaults to [0, 0, 0]. Returns ------- - rbins: Series + rbins : Series The radial bin to which each particle belongs. - bin_edges: ndarray + bin_edges : ndarray Locations of the bin edges. """ @@ -81,7 +83,8 @@ def _bin_particles_by_radius(data: 'SarracenDataFrame', bin_edges = np.logspace(np.log10(r_in), np.log10(r_out), bins+1) else: bin_edges = np.linspace(r_in, r_out, bins+1) - rbins = pd.cut(r, pd.Series(bin_edges)) + + rbins = pd.Series(pd.cut(r, pd.Series(bin_edges))) return rbins, bin_edges @@ -93,7 +96,7 @@ def _get_bin_midpoints(bin_edges: np.ndarray, Parameters ---------- - bin_edges: ndarray + bin_edges : ndarray Locations of the bin edges. log : bool, optional Whether to bin in log scale or not. Defaults to False. diff --git a/sarracen/interpolate/base_backend.py b/sarracen/interpolate/base_backend.py index 0e62bf3..27326d2 100644 --- a/sarracen/interpolate/base_backend.py +++ b/sarracen/interpolate/base_backend.py @@ -1,4 +1,3 @@ -from typing import Tuple from numba.core.registry import CPUDispatcher from numpy import ndarray, zeros @@ -38,7 +37,7 @@ def interpolate_2d_render_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: """ Interpolate 2D vector data to a pair of 2D pixel grids. """ return zeros((y_pixels, x_pixels)), zeros((y_pixels, x_pixels)) @@ -106,7 +105,7 @@ def interpolate_3d_projection_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: """ Interpolate 3D vector data to a pair of 2D pixel grids using column projection.""" return zeros((y_pixels, x_pixels)), zeros((y_pixels, x_pixels)) @@ -147,7 +146,7 @@ def interpolate_3d_cross_vec(x: ndarray, x_min: float, x_max: float, y_min: float, - y_max: float) -> Tuple[ndarray, ndarray]: + y_max: float) -> tuple[ndarray, ndarray]: """ Interpolate 3D vector data to a pair of 2D pixel grids using a 3D cross-section at a specific z value. diff --git a/sarracen/interpolate/cpu_backend.py b/sarracen/interpolate/cpu_backend.py index eb0d2c2..cd449c3 100644 --- a/sarracen/interpolate/cpu_backend.py +++ b/sarracen/interpolate/cpu_backend.py @@ -1,4 +1,3 @@ -from typing import Tuple from numba import njit, prange, get_num_threads from numba.core.registry import CPUDispatcher @@ -47,7 +46,7 @@ def interpolate_2d_render_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: if exact: return (CPUBackend._exact_2d_render(x, y, weight_x, h, x_pixels, y_pixels, x_min, x_max, @@ -133,7 +132,7 @@ def interpolate_3d_projection_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: if exact: return (CPUBackend._exact_3d_project(x, y, weight_x, h, x_pixels, y_pixels, x_min, x_max, @@ -182,7 +181,7 @@ def interpolate_3d_cross_vec(x: ndarray, x_min: float, x_max: float, y_min: float, - y_max: float) -> Tuple[ndarray, ndarray]: + y_max: float) -> tuple[ndarray, ndarray]: return (CPUBackend._fast_2d(x, y, z, z_slice, weight_x, h, weight_function, kernel_radius, x_pixels, y_pixels, x_min, x_max, y_min, y_max, 3), diff --git a/sarracen/interpolate/gpu_backend.py b/sarracen/interpolate/gpu_backend.py index 6247d9e..e1e59b6 100644 --- a/sarracen/interpolate/gpu_backend.py +++ b/sarracen/interpolate/gpu_backend.py @@ -1,5 +1,4 @@ import math -from typing import Tuple import numpy as np from numba import cuda @@ -48,7 +47,7 @@ def interpolate_2d_render_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: if exact: return (GPUBackend._exact_2d_render(x, y, weight_x, h, x_pixels, y_pixels, x_min, x_max, @@ -133,7 +132,7 @@ def interpolate_3d_projection_vec(x: ndarray, x_max: float, y_min: float, y_max: float, - exact: bool) -> Tuple[ndarray, ndarray]: + exact: bool) -> tuple[ndarray, ndarray]: if exact: return (GPUBackend._exact_3d_project(x, y, weight_x, h, x_pixels, y_pixels, x_min, x_max, @@ -182,7 +181,7 @@ def interpolate_3d_cross_vec(x: ndarray, x_min: float, x_max: float, y_min: float, - y_max: float) -> Tuple[ndarray, ndarray]: + y_max: float) -> tuple[ndarray, ndarray]: return (GPUBackend._fast_2d(x, y, z, z_slice, weight_x, h, weight_function, kernel_radius, x_pixels, y_pixels, x_min, x_max, y_min, y_max, 3), diff --git a/sarracen/interpolate/interpolate.py b/sarracen/interpolate/interpolate.py index 796661a..62fabef 100644 --- a/sarracen/interpolate/interpolate.py +++ b/sarracen/interpolate/interpolate.py @@ -9,27 +9,35 @@ from ..interpolate import BaseBackend, CPUBackend, GPUBackend from ..kernels import BaseKernel -from typing import Tuple, Union, Optional, Type, Literal +from typing import TYPE_CHECKING, Literal, TypeAlias +if TYPE_CHECKING: + from ..sarracen_dataframe import SarracenDataFrame -def _default_xy(data: 'SarracenDataFrame', # noqa: F821 - x: Union[str, None], - y: Union[str, None]) -> Tuple[str, str]: + +Bounds: TypeAlias = tuple[float | None, float | None] +VectorLike: TypeAlias = np.ndarray | list | tuple +OriginLike: TypeAlias = VectorLike | pd.Series | str | None + + +def _default_xy(data: 'SarracenDataFrame', + x: str | None, + y: str | None) -> tuple[str, str]: """ Utility function to determine the x & y columns to use during 2D interpolation. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x, y: str + x, y : str or None The x and y directional column labels passed to the interpolation function. Returns ------- - x, y: str + x, y : str The directional column labels to use in interpolation. """ if x is None: @@ -40,25 +48,25 @@ def _default_xy(data: 'SarracenDataFrame', # noqa: F821 return x, y -def _default_xyz(data: 'SarracenDataFrame', # noqa: F821 - x: Union[str, None], - y: Union[str, None], - z: Union[str, None]) -> Tuple[str, str, str]: +def _default_xyz(data: 'SarracenDataFrame', + x: str | None, + y: str | None, + z: str | None) -> tuple[str, str, str]: """ Utility function to determine the x, y and z columns to use during 3-D interpolation. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x, y, z: str + x, y, z : str or None The x, y and z directional column labels passed to the interpolation function. Returns ------- - x, y, z: str + x, y, z : str The directional column labels to use in interpolation. """ xcol = data.xcol @@ -78,32 +86,26 @@ def _default_xyz(data: 'SarracenDataFrame', # noqa: F821 return x, y, z -def _default_bounds(x_data: np.ndarray, - y_data: np.ndarray, - xlim: Union[Tuple[Union[float, None], Union[float, None]], - None], - ylim: Union[Tuple[Union[float, None], Union[float, None]], - None]) -> Tuple[Tuple[float, float], - Tuple[float, float]]: +def _default_bounds(x_data: np.ndarray | pd.Series, + y_data: np.ndarray | pd.Series, + xlim: Bounds | None, + ylim: Bounds | None) -> tuple[tuple[float, float], + tuple[float, float]]: """ Utility function to determine the 2-dimensional boundaries to use in 2D interpolation. Parameters ---------- - data: SarracenDataFrame - The particle dataset to interpolate over. - x, y: str - The directional column labels that will be used in interpolation. - x_data, y_data, z_data: ndarray + x_data, y_data : ndarray The particle coordinate data. - xlim, ylim: tuple of float + xlim, ylim : tuple of float or None The minimum and maximum values passed to the interpolation function, in particle data space. Returns ------- - xlim, ylim: tuple of float + xlim, ylim : tuple of float The minimum and maximum values to use in interpolation, in particle data space. Defaults to the maximum and minimum values of `x` and `y`, snapped to the nearest integer. @@ -122,26 +124,26 @@ def _default_bounds(x_data: np.ndarray, return (x_min, x_max), (y_min, y_max) -def _set_pixels(x_pixels: Union[int, None], - y_pixels: Union[int, None], - xlim: Tuple[float, float], - ylim: Tuple[float, float]) -> Tuple[int, int]: +def _set_pixels(x_pixels: int | None, + y_pixels: int | None, + xlim: tuple[float, float], + ylim: tuple[float, float]) -> tuple[int, int]: """ Utility function to determine the number of pixels to interpolate over in 2D interpolation. Parameters ---------- - x_pixels, y_pixels: int + x_pixels, y_pixels : int The number of pixels in the x & y directions passed to the interpolation function. - xlim, ylim: tuple of float + xlim, ylim : tuple of float The minimum and maximum values to use in interpolation, in particle data space. Returns ------- - x_pixels, y_pixels: int + x_pixels, y_pixels : int The number of pixels in the x & y directions to use in 2D interpolation. """ @@ -161,7 +163,7 @@ def _set_pixels(x_pixels: Union[int, None], return x_pixels, y_pixels -def _verify_columns(data: 'SarracenDataFrame', # noqa: F821 +def _verify_columns(data: 'SarracenDataFrame', x: str, y: str) -> None: """ @@ -169,9 +171,9 @@ def _verify_columns(data: 'SarracenDataFrame', # noqa: F821 Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x, y: str + x, y : str The directional column labels that will be used in interpolation. Raises @@ -199,18 +201,18 @@ def _verify_columns(data: 'SarracenDataFrame', # noqa: F821 def _check_boundaries(x_pixels: int, y_pixels: int, - xlim: Tuple[float, float], - ylim: Tuple[float, float]) -> None: + xlim: tuple[float, float], + ylim: tuple[float, float]) -> None: """ Verify that the pixel count and boundaries of a 2D plot describe a valid region. Parameters ---------- - x_pixels, y_pixels: int + x_pixels, y_pixels : int The number of pixels in the x & y directions passed to the interpolation function. - xlim, ylim: tuple of float + xlim, ylim : tuple of float The minimum and maximum values to use in interpolation, in particle data space. @@ -231,7 +233,7 @@ def _check_boundaries(x_pixels: int, raise ValueError("`y_pixels` must be greater than zero!") -def _check_dimension(data: 'SarracenDataFrame', # noqa: F821 +def _check_dimension(data: 'SarracenDataFrame', dim: Literal[2, 3]) -> None: """ Verify that a given dataset describes data with a required number of @@ -239,9 +241,9 @@ def _check_dimension(data: 'SarracenDataFrame', # noqa: F821 Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - dim: [2, 3] + dim : [2, 3] The number of required dimensions. Raises @@ -255,28 +257,26 @@ def _check_dimension(data: 'SarracenDataFrame', # noqa: F821 raise ValueError(f"Dataset is not {dim}-dimensional.") -def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 +def _rotate_data(data: 'SarracenDataFrame', x_data: np.ndarray, y_data: np.ndarray, z_data: np.ndarray, - rotation: Union[np.ndarray, list, Rotation, None], - rot_origin: Union[np.ndarray, list, pd.Series, - str, None]) -> Tuple[np.ndarray, - np.ndarray, - np.ndarray]: + rotation: VectorLike | Rotation | None, + rot_origin: OriginLike) -> tuple[np.ndarray, np.ndarray, + np.ndarray]: """ Rotate vector data in a particle dataset. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x_data, y_data, z_data: ndarray + x_data, y_data, z_data : ndarray The directional vector data. - rotation: array_like or SciPy Rotation + rotation : array_like or SciPy Rotation The rotation to apply to the vector data. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'] + rot_origin : array_like or ['com', 'midpoint'] Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the @@ -285,7 +285,7 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 Returns ------- - x_data, y_data, z_data: ndarray + x_data, y_data, z_data : ndarray The rotated x, y, and z directional data. """ vectors = np.array([x_data, y_data, z_data]).transpose() @@ -313,7 +313,7 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 "'com' or 'midpoint'") # user specified point - elif isinstance(rot_origin, (list, pd.Series, np.ndarray)): + elif isinstance(rot_origin, (list, tuple, pd.Series, np.ndarray)): if len(rot_origin) != 3: raise ValueError("rot_origin should specify [x, y, z] point.") rot_origin_arr = np.asarray(rot_origin, dtype=x_data.dtype) @@ -329,15 +329,13 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 return vectors[:, 0], vectors[:, 1], vectors[:, 2] -def _rotate_xyz(data: 'SarracenDataFrame', # noqa: F821 +def _rotate_xyz(data: 'SarracenDataFrame', x: str, y: str, z: str, - rotation: Union[np.ndarray, list, Rotation, None], - rot_origin: Union[np.ndarray, list, - str, None]) -> Tuple[np.ndarray, - np.ndarray, - np.ndarray]: + rotation: VectorLike | Rotation | None, + rot_origin: OriginLike) -> tuple[np.ndarray, np.ndarray, + np.ndarray]: """ Rotate positional data in a particle dataset. @@ -348,14 +346,14 @@ def _rotate_xyz(data: 'SarracenDataFrame', # noqa: F821 Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x, y, z: str + x, y, z : str Directional column labels containing the positional column labels. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the @@ -364,49 +362,48 @@ def _rotate_xyz(data: 'SarracenDataFrame', # noqa: F821 Returns ------- - x_data, y_data, z_data: ndarray + x_data, y_data, z_data : ndarray The rotated x, y, and z directional data. """ - rotated_x, rotated_y, rotated_z = _rotate_data(data, data[data.xcol], - data[data.ycol], - data[data.zcol], rotation, + rotated_x, rotated_y, rotated_z = _rotate_data(data, + data[data.xcol].to_numpy(), + data[data.ycol].to_numpy(), + data[data.zcol].to_numpy(), + rotation, rot_origin) x_data = rotated_x if x == data.xcol else \ rotated_y if x == data.ycol else \ - rotated_z if x == data.zcol else data[x] + rotated_z if x == data.zcol else data[x].to_numpy() y_data = rotated_x if y == data.xcol else \ rotated_y if y == data.ycol else \ - rotated_z if y == data.zcol else data[y] + rotated_z if y == data.zcol else data[y].to_numpy() z_data = rotated_x if z == data.xcol else \ rotated_y if z == data.ycol else \ - rotated_z if z == data.zcol else data[z] + rotated_z if z == data.zcol else data[z].to_numpy() return x_data, y_data, z_data -def _corotate(corotation: Union[np.ndarray, list], - rotation: Union[np.ndarray, list, - Rotation, None]) -> Tuple[Union[np.ndarray, - list, - Rotation], - Union[np.ndarray, - list]]: +def _corotate(corotation: VectorLike, + rotation: VectorLike | Rotation | None) -> tuple[VectorLike | + Rotation, + VectorLike]: """ Calculates the rotation matrix for a corotating frame. Parameters ---------- - corotation: array_like + corotation : array_like The x, y, z coordinates of two locations which determines the corotating frame. Each coordinate is also array_like. - rotation: array_like or SciPy Rotation + rotation : array_like or SciPy Rotation An additional rotation to apply to the corotating frame. Returns ------- - rotation: array_like or SciPy Rotation + rotation : array_like or SciPy Rotation The rotation to apply to the data before interpolation. - rot_origin: array_like + rot_origin : array_like Point of rotation of the data. """ corotation[1][0] -= corotation[0][0] @@ -428,15 +425,14 @@ def _corotate(corotation: Union[np.ndarray, list], return rotation, rot_origin -def _get_mass(data: 'SarracenDataFrame') -> Union[np.ndarray, # noqa: F821 - float]: +def _get_mass(data: 'SarracenDataFrame') -> np.ndarray | float: if data.mcol is None: return data.params['mass'] return data[data.mcol].to_numpy() -def _get_density(data: 'SarracenDataFrame') -> np.ndarray: # noqa: F821 +def _get_density(data: 'SarracenDataFrame') -> np.ndarray: if data.rhocol is None: hfact = data.params['hfact'] mass = _get_mass(data) @@ -445,8 +441,8 @@ def _get_density(data: 'SarracenDataFrame') -> np.ndarray: # noqa: F821 return data[data.rhocol].to_numpy() -def _get_weight(data: 'SarracenDataFrame', # noqa: F821 - target: Union[str, np.ndarray], +def _get_weight(data: 'SarracenDataFrame', + target: str | np.ndarray, dens_weight: bool) -> np.ndarray: if type(target) is str: @@ -472,12 +468,12 @@ def _get_weight(data: 'SarracenDataFrame', # noqa: F821 return target_data * mass_data / rho_data -def _get_smoothing_lengths(data: 'SarracenDataFrame', # noqa: F821 +def _get_smoothing_lengths(data: 'SarracenDataFrame', hmin: bool, x_pixels: int, y_pixels: int, - xlim: Tuple[float, float], - ylim: Tuple[float, float]) -> np.ndarray: + xlim: tuple[float, float], + ylim: tuple[float, float]) -> np.ndarray: """ Return smoothing lengths, imposing a min length if hmin is True. """ if hmin: @@ -490,19 +486,17 @@ def _get_smoothing_lengths(data: 'SarracenDataFrame', # noqa: F821 return h_data -def interpolate_2d(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_2d(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, + x: str | None = None, + y: str | None = None, + kernel: BaseKernel | None = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, hmin: bool = False) -> np.ndarray: @@ -518,32 +512,34 @@ def interpolate_2d(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target smoothing data. - x, y: str, optional + x, y : str, optional Column labels of the directional axes. Defaults to the x & y columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - exact: bool, optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + exact : bool, optional Whether to use exact interpolation of the data. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). @@ -596,23 +592,21 @@ def interpolate_2d(data: 'SarracenDataFrame', # noqa: F821 return grid -def interpolate_2d_vec(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_2d_vec(data: 'SarracenDataFrame', target_x: str, target_y: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, + x: str | None = None, + y: str | None = None, + kernel: BaseKernel | None = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, - hmin: bool = False) -> Tuple[np.ndarray, np.ndarray]: + hmin: bool = False) -> tuple[np.ndarray, np.ndarray]: """ Interpolate vector particle data across two directional axes to a 2D grid of particles. @@ -625,39 +619,41 @@ def interpolate_2d_vec(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target_x, target_y: str + target_x, target_y : str Column labels of the target vector. - x, y: str, optional + x, y : str, optional Column labels of the directional axes. Defaults to the x & y columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - exact: bool, optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + exact : bool, optional Whether to use exact interpolation of the data. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). Returns ------- - output_x, output_y: ndarray (2-Dimensional) + output_x, output_y : ndarray (2-Dimensional) The interpolated output images, in a 2-dimensional numpy arrays. Dimensions are structured in reverse order, where (x, y) -> [y, x]. @@ -708,15 +704,15 @@ def interpolate_2d_vec(data: 'SarracenDataFrame', # noqa: F821 return (gridx, gridy) -def interpolate_2d_line(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_2d_line(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - backend: Union[str, None] = None, + x: str | None = None, + y: str | None = None, + kernel: BaseKernel | None = None, + pixels: int | None = None, + xlim: Bounds | float | int | None = None, + ylim: Bounds | float | int | None = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, hmin: bool = False) -> np.ndarray: @@ -733,28 +729,32 @@ def interpolate_2d_line(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target smoothing data. - x, y: str, optional + x, y : str, optional Column labels of the directional axes. Defaults to the x & y columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - pixels: int, optional + pixels : int, optional Number of points in the resulting line plot in the x-direction. - xlim, ylim: tuple of float, optional - Starting and ending coordinates of the cross-section line (in particle - data space). Defaults to the minimum and maximum values of `x` and `y`. - backend: ['cpu', 'gpu'], optional + xlim, ylim : int, float, or tuple of float or None, optional + Coordinate bounds of the cross-sectional line. A scalar value is + treated as both bounds for that axis, fixing the line along that + constant coordinate value. A tuple gives the upper and lower bounds. + Either value of the tuple may be `None`, which defaults to the minimum + or maximum value of `x` and `y`, respectively. If omitted, the limit + defaults to the minimum and maximum values. + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). @@ -820,17 +820,17 @@ def interpolate_2d_line(data: 'SarracenDataFrame', # noqa: F821 return grid -def interpolate_3d_line(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_line(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - zlim: Union[Tuple[float, float], None] = None, - backend: Union[str, None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, + pixels: int | None = None, + xlim: Bounds | float | int | None = None, + ylim: Bounds | float | int | None = None, + zlim: Bounds | float | int | None = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, hmin: bool = False) -> np.ndarray: @@ -846,37 +846,40 @@ def interpolate_3d_line(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target variable. - x, y, z: str, optional + x, y, z : str, optional Column labels of the directional axes. Defaults to the x, y & z columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - pixels: int, optional + pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim, zlim: tuple of float, optional - Starting and ending coordinates of the cross-section line (in particle - data space). Defaults to the minimum and maximum values of `x`, `y`, - and `z`. - backend: ['cpu', 'gpu'], optional + xlim, ylim, zlim : float, int, or tuple of float or None, optional + Coordinate bounds of the cross-sectional line. A scalar value is + treated as both bounds for that axis, fixing the line along that + constant coordinate value. A tuple gives the upper and lower bounds. + Either value of the tuple may be `None`, which defaults to the minimum + or maximum value of `x`, `y`, and `z`, respectively. If omitted, the + limit defaults to the minimum and maximum values. + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). Returns ------- - output: ndarray (1-Dimensional) + output : ndarray (1-Dimensional) The interpolated output line. Raises @@ -906,6 +909,9 @@ def interpolate_3d_line(data: 'SarracenDataFrame', # noqa: F821 z2 = data.loc[:, z].min() if zlim is None or zlim[1] is None else zlim[1] zlim = z1, z2 + if zlim is None or zlim[0] is None or zlim[1] is None: + raise ValueError('Zero area cross section!') + xlim, ylim = _default_bounds(data[x], data[y], xlim, ylim) if ylim[1] == ylim[0] and xlim[1] == xlim[0] and zlim[1] == zlim[0]: raise ValueError('Zero length cross section!') @@ -943,23 +949,22 @@ def interpolate_3d_line(data: 'SarracenDataFrame', # noqa: F821 return grid -def interpolate_3d_proj(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_proj(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, + x: str | None = None, + y: str | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - corotation: Union[np.ndarray, list, None] = None, - rotation: Union[np.ndarray, list, - Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, + corotation: VectorLike | None = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, exact: bool = False, - backend: Union[str, None] = None, - dens_weight: Union[bool, None] = None, + backend: str | None = None, + dens_weight: bool | None = None, normalize: bool = True, hmin: bool = False) -> np.ndarray: """ @@ -973,48 +978,50 @@ def interpolate_3d_proj(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target smoothing data. - x, y: str, optional + x, y : str, optional Column labels of the directional axes. Defaults to the x & y columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - integral_samples: int, optional + integral_samples : int, optional Number of sample points to take when approximating the 2D column kernel. - corotation: array_like, optional + corotation : array_like, optional The x, y, z coordinates of two locations which determines the corotating frame. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - exact: bool, optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + exact : bool, optional Whether to use exact interpolation of the data. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to True for column-integrated views, when the target is not density, and False for everything else. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). @@ -1083,28 +1090,25 @@ def interpolate_3d_proj(data: 'SarracenDataFrame', # noqa: F821 return grid -def interpolate_3d_vec(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_vec(data: 'SarracenDataFrame', target_x: str, target_y: str, target_z: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, + x: str | None = None, + y: str | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, - Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, - hmin: bool = False) -> Tuple[np.ndarray, np.ndarray]: + hmin: bool = False) -> tuple[np.ndarray, np.ndarray]: """ Interpolate 3D vector particle data to a 2D grid of pixels. @@ -1116,50 +1120,52 @@ def interpolate_3d_vec(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target_x, target_y, target_z: str + target_x, target_y, target_z : str Column labels of the target vector. - x, y: str, optional + x, y : str, optional Column labels of the directional axes. Defaults to the x & y columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - integral_samples: int, optional + integral_samples : int, optional Number of sample points to take when approximating the 2D column kernel. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - exact: bool, optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + exact : bool, optional Whether to use exact interpolation of the data. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). Returns ------- - output_x, output_y: ndarray (2-Dimensional) + output_x, output_y : ndarray (2-Dimensional) The interpolated output images. Dimensions are structured in reverse order, where (x, y) -> [y, x]. @@ -1192,9 +1198,11 @@ def interpolate_3d_vec(data: 'SarracenDataFrame', # noqa: F821 if target_z not in data.columns: raise KeyError(f"z-directional target column '{target_z}' does not " f"exist in the provided dataset.") - target_x_data, target_y_data, _ = _rotate_data(data, data[target_x], - data[target_y], - data[target_z], rotation, + target_x_data, target_y_data, _ = _rotate_data(data, + data[target_x].to_numpy(), + data[target_y].to_numpy(), + data[target_z].to_numpy(), + rotation, rot_origin) wx_data = _get_weight(data, target_x_data, dens_weight) @@ -1226,22 +1234,21 @@ def interpolate_3d_vec(data: 'SarracenDataFrame', # noqa: F821 return (gridx, gridy) -def interpolate_3d_cross(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_cross(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - z_slice: Union[float, None] = None, - kernel: Union[BaseKernel, None] = None, - corotation: Union[np.ndarray, list, None] = None, - rotation: Union[np.ndarray, list, - Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - backend: Union[str, None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + z_slice: float | None = None, + kernel: BaseKernel | None = None, + corotation: VectorLike | None = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, hmin: bool = False) -> np.ndarray: @@ -1257,45 +1264,47 @@ def interpolate_3d_cross(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame The particle data to interpolate over. - target: str + target : str The column label of the target smoothing data. - x, y, z: str, optional + x, y, z : str, optional The column labels of the directional data to interpolate over. Defaults to the x, y, and z columns detected in `data`. - z_slice: float, optional + z_slice : float, optional The z-axis value to take the cross-section at. Defaults to the midpoint of the z-directional data. - kernel: BaseKernel, optional + kernel : BaseKernel, optional The kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - corotation: array_like, optional + corotation : array_like, optional The x, y, z coordinates of two locations which determines the corotating frame. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - backend: ['cpu', 'gpu'], optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). @@ -1361,29 +1370,25 @@ def interpolate_3d_cross(data: 'SarracenDataFrame', # noqa: F821 return grid -def interpolate_3d_cross_vec(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_cross_vec(data: 'SarracenDataFrame', target_x: str, target_y: str, target_z: str, - z_slice: Union[float, None] = None, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - rotation: Union[np.ndarray, list, Rotation, - None] = None, - rot_origin: Union[np.ndarray, list, - str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - backend: Union[str, None] = None, + z_slice: float | None = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, - hmin: bool = False) -> Tuple[np.ndarray, + hmin: bool = False) -> tuple[np.ndarray, np.ndarray]: """ Interpolate 3D vector particle data to a 2D grid, using a 3D cross-section. @@ -1397,48 +1402,50 @@ def interpolate_3d_cross_vec(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame The particle data to interpolate over. - target_x, target_y, target_z: str + target_x, target_y, target_z : str The column labels of the target vector. - z_slice: float, optional + z_slice : float, optional The z-axis value to take the cross-section at. Defaults to the midpoint of the z-directional data. - x, y, z: str, optional + x, y, z : str, optional The column labels of the directional data to interpolate over. Defaults to the x, y, and z columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional The kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: float, optional + xlim, ylim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x` and `y`. - backend: ['cpu', 'gpu'], optional + `None` in a tuple defaults to the minimum or maximum value, + respectively. + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). Returns ------- - output_x, output_y: ndarray (2-Dimensional) + output_x, output_y : ndarray (2-Dimensional) The interpolated output images. Dimensions are structured in reverse order, where (x, y) -> [y, x]. @@ -1469,9 +1476,11 @@ def interpolate_3d_cross_vec(data: 'SarracenDataFrame', # noqa: F821 x_pixels, y_pixels = _set_pixels(x_pixels, y_pixels, xlim, ylim) _check_boundaries(x_pixels, y_pixels, xlim, ylim) - target_x_data, target_y_data, _ = _rotate_data(data, data[target_x], - data[target_y], - data[target_z], rotation, + target_x_data, target_y_data, _ = _rotate_data(data, + data[target_x].to_numpy(), + data[target_y].to_numpy(), + data[target_z].to_numpy(), + rotation, rot_origin) wx_data = _get_weight(data, target_x_data, dens_weight) @@ -1501,24 +1510,21 @@ def interpolate_3d_cross_vec(data: 'SarracenDataFrame', # noqa: F821 return (gridx, gridy) -def interpolate_3d_grid(data: 'SarracenDataFrame', # noqa: F821 +def interpolate_3d_grid(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - rotation: Union[np.ndarray, list, - Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - z_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - zlim: Union[Tuple[float, float], None] = None, - backend: Union[str, None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + z_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + zlim: Bounds | None = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = True, hmin: bool = False) -> np.ndarray: @@ -1533,39 +1539,40 @@ def interpolate_3d_grid(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame The particle data to interpolate over. - target: str + target : str The column label of the target data. - x, y, z: str, optional + x, y, z : str, optional The column labels of the directional data to interpolate over. Defaults to the x, y, and z columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional The kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels, z_pixels: int, optional + x_pixels, y_pixels, z_pixels : int, optional Number of pixels in the output image in the x, y & z directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim, zlim: tuple of float, optional + xlim, ylim, zlim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x`, `y` - and `z`. - backend: ['cpu', 'gpu'], optional + and `z`. `None` in a tuple defaults to the minimum or maximum value, + respectively. + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, the target will be multiplied by density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). @@ -1599,7 +1606,13 @@ def interpolate_3d_grid(data: 'SarracenDataFrame', # noqa: F821 if not ylim: ylim = (None, None) xlim, ylim = _default_bounds(x_data, y_data, xlim, ylim) - zlim = zlim if zlim else (z_data.min(), z_data.max()) + + z1 = data.loc[:, z].min() if zlim is None or zlim[0] is None else zlim[0] + z2 = data.loc[:, z].max() if zlim is None or zlim[1] is None else zlim[1] + zlim = z1, z2 + + if zlim is None or zlim[0] is None or zlim[1] is None: + raise ValueError('Zero area cross section!') x_pixels, y_pixels = _set_pixels(x_pixels, y_pixels, xlim, ylim) if z_pixels is None: @@ -1637,19 +1650,19 @@ def interpolate_3d_grid(data: 'SarracenDataFrame', # noqa: F821 return grid -def get_backend(code: str) -> Type[BaseBackend]: +def get_backend(code: str) -> type[BaseBackend]: """ Get the interpolation backend associated with a string code. Parameters ---------- - code: str + code : str The code associated with the particular backend. At the moment, 'cpu' for the CPU backend, and 'gpu' for the GPU backend are supported. Returns ------- - BaseBackend: The backend to use for interpolation. + BaseBackend : The backend to use for interpolation. """ if code == 'cpu': return CPUBackend diff --git a/sarracen/kernels/base_kernel.py b/sarracen/kernels/base_kernel.py index 6d82f4a..7a38172 100644 --- a/sarracen/kernels/base_kernel.py +++ b/sarracen/kernels/base_kernel.py @@ -1,8 +1,8 @@ +from collections.abc import Callable import math import numpy as np from numba import njit, prange -from typing import Callable class BaseKernel: @@ -41,7 +41,7 @@ def get_column_kernel(self, samples: int = 1000) -> np.ndarray: Parameters ---------- - samples: int + samples : int Number of sample points to calculate when approximating the kernel. Returns @@ -74,7 +74,7 @@ def get_column_kernel_func(self, samples: int) -> Callable[[float, int], Parameters ---------- - samples: int + samples : int Number of sample points to calculate when approximating the kernel. Returns diff --git a/sarracen/kernels/cubic_spline_exact.py b/sarracen/kernels/cubic_spline_exact.py index c231584..0148015 100644 --- a/sarracen/kernels/cubic_spline_exact.py +++ b/sarracen/kernels/cubic_spline_exact.py @@ -1,7 +1,6 @@ import math from numba import njit -from typing import Tuple @njit @@ -12,11 +11,11 @@ def line_int(r0: float, d1: float, d2: float, h: float) -> float: Parameters ---------- - r0: float + r0 : float Distance between the contributing particle and the line. - d1, d2: float + d1, d2 : float Distance from the endpoint of `r0` to each endpoint of the line. - h: float + h : float Smoothing length of the contributing particle. Returns @@ -63,10 +62,10 @@ def _full_2d_mod(phi: float, q0: float) -> float: Parameters ---------- - phi: float + phi : float Angle between `q0` and the endpoint of the line, relative to the contributing particle. - q0: float + q0 : float The distance between the contributing particle and the line, scaled by the smoothing length of the particle. @@ -131,10 +130,10 @@ def _f1_2d(phi: float, q0: float) -> float: Parameters ---------- - phi: float + phi : float Angle between `q0` and the endpoint of the line segment, relative to the contributing particle. - q0: float + q0 : float The distance between the contributing particle and the line, scaled by the smoothing length of the particle. @@ -164,10 +163,10 @@ def _f2_2d(phi: float, q0: float) -> float: Parameters ---------- - phi: float + phi : float Angle between `q0` and the endpoint of the line segment, relative to the contributing particle. - q0: float + q0 : float The distance between the contributing particle and the line, scaled by the smoothing length of the particle. @@ -205,7 +204,7 @@ def _f3_2d(phi: float) -> float: Parameters ---------- - phi: float + phi : float Angle pointing towards the endpoint of the line segment, relative to the contributing particle. @@ -233,13 +232,13 @@ def surface_int(r0: float, Parameters ---------- - r0: float + r0 : float Distance between the contributing particle and the target surface. - x1, y1, x2, y2: float + x1, y1, x2, y2 : float Upper and lower bounds of the target surface, relative to r0. - wx, wy: float + wx, wy : float The size of a single pixel. - h: float + h : float The smoothing length of the contributing particle. Returns @@ -289,13 +288,13 @@ def _line_int3d(r0: float, r1: float, d1: float, d2: float, h: float) -> float: Parameters ---------- - r0: float + r0 : float Distance between the contributing particle and the target surface. - r1: float + r1 : float Distance between the endpoint of `r0` and the line. - d1, d2: float + d1, d2 : float Distance from the endpoint of `r1` to each endpoint of the line. - h: float + h : float Smoothing length of the contributing particle. Returns @@ -359,13 +358,13 @@ def _full_integral_3d(d: float, r0: float, r1: float, h: float) -> float: Parameters ---------- - d: float + d : float Distance of the line. - r0: float + r0 : float Distance between the target surface and the contributing particle. - r1: float + r1 : float Distance between the end of `r0` and the line. - h: float + h : float The smoothing length of the contributing particle. Returns @@ -457,7 +456,7 @@ def _full_integral_3d(d: float, r0: float, r1: float, h: float) -> float: @njit def get_I_terms(cosp: float, a2: float, - a: float) -> Tuple[float, float, float, float, float, float]: + a: float) -> tuple[float, float, float, float, float, float]: """Calculate I constants for calculations in _full_integral_3d""" cosp2 = cosp * cosp p = math.acos(cosp) diff --git a/sarracen/readers/read_gasoline.py b/sarracen/readers/read_gasoline.py index a3c2005..97c000f 100644 --- a/sarracen/readers/read_gasoline.py +++ b/sarracen/readers/read_gasoline.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Literal, TypeAlias, overload import numpy as np import pandas as pd @@ -8,12 +8,26 @@ from ..sarracen_dataframe import SarracenDataFrame +GasolineFrames: TypeAlias = tuple[SarracenDataFrame, SarracenDataFrame, + SarracenDataFrame] +GasolineDicts: TypeAlias = tuple[dict, dict, dict, dict] + + +@overload +def read_gasoline(filename: str, + outtype: Literal["sarracen"] = "sarracen" + ) -> GasolineFrames | int: ... +@overload # noqa: E302 def read_gasoline(filename: str, - outtype: str = "sarracen") -> Union[Tuple[SarracenDataFrame, - SarracenDataFrame, - SarracenDataFrame], - Tuple[dict, dict, - dict, dict], int]: + outtype: Literal["dic", "dict", "dictionary"] + ) -> GasolineDicts | int: ... +@overload # noqa: E302 +def read_gasoline(filename: str, + outtype: str + ) -> GasolineFrames | GasolineDicts | int: ... +def read_gasoline(filename: str, # noqa: E302 + outtype: str = "sarracen" + ) -> GasolineFrames | GasolineDicts | int: """ Read data from a Gasoline tipsy file. @@ -34,10 +48,15 @@ def read_gasoline(filename: str, Returns ------- - list of SarracenDataFrames or Python dictionaries + tuple of SarracenDataFrame or dict + + Raises + ------ + ValueError + If the header and file size are inconsistent. Notes - -------- + ----- Adapted from `PyTipsy `_. Examples @@ -98,14 +117,14 @@ def read_gasoline(filename: str, fp.read(4) # File is borked if this is true elif (fs != 28+48*ng+36*nd+44*ns): - print("Tipsy ERROR: Header and file size inconsistend") + print("Tipsy ERROR: Header and file size inconsistent") print("Estimates: Header bytes: 28 or 32 (either is OK)") print(" ngas: ", ng, " bytes:", 48*ng) print(" ndark: ", nd, " bytes:", 38*nd) print(" nstar: ", ns, " bytes:", 44*ns) print("Actual File bytes:", fs, " does not work") fp.close() - return 1 + raise ValueError("Header and file size inconsistent") # Make dicitonaries for data catg = {'mass': np.zeros(ng), 'pos': np.zeros((ng, 3)), diff --git a/sarracen/readers/read_gradsph.py b/sarracen/readers/read_gradsph.py index 9d110ed..dbb3ceb 100644 --- a/sarracen/readers/read_gradsph.py +++ b/sarracen/readers/read_gradsph.py @@ -1,14 +1,11 @@ -from typing import Union, List - import pandas as pd from ..sarracen_dataframe import SarracenDataFrame def read_gradsph(filename: str, - separate_types: str = 'sinks') -> Union[List[ - SarracenDataFrame], - SarracenDataFrame]: + separate_types: str | None = 'sinks' + ) -> list[SarracenDataFrame] | SarracenDataFrame: """ Read data from a GradSPH dump file. diff --git a/sarracen/readers/read_marisa.py b/sarracen/readers/read_marisa.py index 6c13a95..b647690 100644 --- a/sarracen/readers/read_marisa.py +++ b/sarracen/readers/read_marisa.py @@ -1,4 +1,4 @@ -from typing import Tuple, IO, Type +from typing import IO import numpy as np import pandas as pd @@ -114,7 +114,7 @@ class MARISAIO_TAGS(IntEnum): # 2: sets the reference point at the end of the file -def _marisa_read_capture_pattern(fp: IO) -> Tuple[int, int]: +def _marisa_read_capture_pattern(fp: IO) -> tuple[int, int]: if fp.read(7) != b"marisa\x00": raise AssertionError("Capture pattern not present. " "Is this a valid data file?") @@ -128,13 +128,13 @@ def _marisa_read_tag(fp: IO) -> int: return tag -def _marisa_read_data(fp: IO, dtype: Type[np.number]) -> np.ndarray: +def _marisa_read_data(fp: IO, dtype: type[np.number]) -> np.ndarray: size = int.from_bytes(fp.read(4), byteorder='little') data = fp.read(size) return np.frombuffer(data, dtype=dtype) -def _marisa_parse_tags(fp: IO) -> Tuple[np.ndarray, np.ndarray]: +def _marisa_parse_tags(fp: IO) -> tuple[np.ndarray, np.ndarray]: currentpos = fp.tell() fp.seek(0, 2) diff --git a/sarracen/readers/read_phantom.py b/sarracen/readers/read_phantom.py index 135f8eb..a76deab 100644 --- a/sarracen/readers/read_phantom.py +++ b/sarracen/readers/read_phantom.py @@ -1,4 +1,4 @@ -from typing import IO, Tuple, Type, Union, List, overload, Literal +from typing import IO, Literal, TypeAlias, overload import numpy as np import pandas as pd @@ -6,6 +6,9 @@ from ..sarracen_dataframe import SarracenDataFrame +PhantomFrames: TypeAlias = list[SarracenDataFrame] | SarracenDataFrame + + def _read_fortran_block(fp: IO, bytesize: int) -> bytes: """ Helper function to read Fortran-written data. @@ -24,15 +27,15 @@ def _read_fortran_block(fp: IO, bytesize: int) -> bytes: return data -def _read_capture_pattern(fp: IO) -> Tuple[Type[np.generic], - Type[np.generic], +def _read_capture_pattern(fp: IO) -> tuple[type[np.generic], + type[np.generic], int, bool]: """ Phantom dump validation plus default real and int sizes.""" start_tag = fp.read(4) # 4-byte Fortran tag - def_types: List[Tuple[Type[np.generic], - Type[np.generic]]] = [(np.int32, np.float64), + def_types: list[tuple[type[np.generic], + type[np.generic]]] = [(np.int32, np.float64), (np.int32, np.float32), (np.int64, np.float64), (np.int64, np.float32)] @@ -121,8 +124,8 @@ def _rename_duplicates(keys: list) -> list: def _read_global_header_block(fp: IO, - dtype: Type[np.generic], - swap_endian: bool) -> Tuple[list, list]: + dtype: type[np.generic], + swap_endian: bool) -> tuple[list, list]: nvars = np.frombuffer(_read_fortran_block(fp, 4), dtype=np.int32)[0] if swap_endian: nvars = nvars.byteswap() @@ -145,8 +148,8 @@ def _read_global_header_block(fp: IO, def _read_global_header(fp: IO, - def_int_dtype: Type[np.generic], - def_real_dtype: Type[np.generic], + def_int_dtype: type[np.generic], + def_real_dtype: type[np.generic], swap_endian: bool) -> dict: """ Read global variables. """ @@ -174,8 +177,8 @@ def _read_array_block(fp: IO, df: pd.DataFrame, n: int, nums: np.ndarray, - def_int_dtype: Type[np.generic], - def_real_dtype: Type[np.generic], + def_int_dtype: type[np.generic], + def_real_dtype: type[np.generic], swap_endian: bool) -> pd.DataFrame: dtypes = [def_int_dtype, np.int8, np.int16, np.int32, np.int64, @@ -204,10 +207,10 @@ def _read_array_block(fp: IO, def _read_array_blocks(fp: IO, - def_int_dtype: Type[np.generic], - def_real_dtype: Type[np.generic], + def_int_dtype: type[np.generic], + def_real_dtype: type[np.generic], mpi_blocks: int, - swap_endian: bool) -> Tuple[pd.DataFrame, pd.DataFrame]: + swap_endian: bool) -> tuple[pd.DataFrame, pd.DataFrame]: """ Read particle data. Block 2 is always for sink particles. The number of MPI blocks is given by @@ -225,8 +228,8 @@ def _read_array_blocks(fp: IO, for j in range(0, mpi_blocks): - n: List[int] = [] - nums: List[np.ndarray] = [] + n: list[int] = [] + nums: list[np.ndarray] = [] for i in range(0, nblocks): start_tag = fp.read(4) @@ -301,20 +304,14 @@ def read_phantom(filename: str, @overload # noqa: E302 def read_phantom(filename: str, separate_types: Literal['sinks'] = 'sinks', - ignore_inactive: bool = True) -> Union[List[ - SarracenDataFrame], - SarracenDataFrame]: ... + ignore_inactive: bool = True) -> PhantomFrames: ... @overload # noqa: E302 def read_phantom(filename: str, separate_types: Literal['all'] = 'all', - ignore_inactive: bool = True) -> Union[List[ - SarracenDataFrame], - SarracenDataFrame]: ... + ignore_inactive: bool = True) -> PhantomFrames: ... def read_phantom(filename: str, # noqa: E302 - separate_types: Union[str, None] = 'sinks', - ignore_inactive: bool = True) -> Union[List[ - SarracenDataFrame], - SarracenDataFrame]: + separate_types: str | None = 'sinks', + ignore_inactive: bool = True) -> PhantomFrames: """ Read data from a Phantom dump file. diff --git a/sarracen/readers/read_phantom_ev.py b/sarracen/readers/read_phantom_ev.py index e709034..29844b0 100644 --- a/sarracen/readers/read_phantom_ev.py +++ b/sarracen/readers/read_phantom_ev.py @@ -1,9 +1,9 @@ -from typing import IO, List, Union +from typing import IO import pandas as pd import numpy as np -def _determine_column_labels(file: IO) -> List[str]: +def _determine_column_labels(file: IO) -> list[str]: """ Determine the column labels. @@ -42,7 +42,7 @@ def _determine_column_labels(file: IO) -> List[str]: return labels -def _infer_type(s: str) -> Union[np.float64, np.int32, str]: +def _infer_type(s: str) -> np.float64 | np.int32 | str: """ Given a value from the ev, determine its type. diff --git a/sarracen/readers/read_shamrock.py b/sarracen/readers/read_shamrock.py index 0974ca2..a72670b 100644 --- a/sarracen/readers/read_shamrock.py +++ b/sarracen/readers/read_shamrock.py @@ -1,4 +1,4 @@ -from typing import IO, List, Dict, Any +from typing import IO, Any import struct import json @@ -95,7 +95,7 @@ def read_from_position(self, offset: int, num_bytes: int) -> bytes: return data -def decode_bytes_to_doubles(byte_data: bytes) -> List[float]: +def decode_bytes_to_doubles(byte_data: bytes) -> list[float]: """ Decodes a byte array into a list of double precision floats. @@ -106,7 +106,7 @@ def decode_bytes_to_doubles(byte_data: bytes) -> List[float]: Returns ------- - List[float] + list of float A list of double-precision floating-point numbers decoded from the byte array. @@ -136,7 +136,7 @@ def get_head_inc(off: int) -> int: return off -def decode_patchdata(pdat: bytes, pdat_layout: List[Dict[str, Any]]) -> dict: +def decode_patchdata(pdat: bytes, pdat_layout: list[dict[str, Any]]) -> dict: """ Decode a patchdata bytearray into a dictionary of numpy arrays. diff --git a/sarracen/render.py b/sarracen/render.py index 609a430..3d6d044 100644 --- a/sarracen/render.py +++ b/sarracen/render.py @@ -9,7 +9,7 @@ data.render_2d(target) """ -from typing import Any, Union, Tuple +from typing import TYPE_CHECKING, Any, TypeAlias import numpy as np import pandas as pd @@ -19,28 +19,36 @@ from matplotlib.axes import Axes from matplotlib.colors import Colormap, LogNorm, SymLogNorm +if TYPE_CHECKING: + from .sarracen_dataframe import SarracenDataFrame + from .interpolate import interpolate_2d_line, interpolate_2d, \ interpolate_3d_proj, interpolate_3d_cross, interpolate_3d_vec, \ interpolate_3d_cross_vec, interpolate_2d_vec, interpolate_3d_line from .kernels import BaseKernel -def _default_axes(data: 'SarracenDataFrame', # noqa: F821 - x: Union[str, None], - y: Union[str, None]) -> Tuple[str, str]: +Bounds: TypeAlias = tuple[float | None, float | None] +VectorLike: TypeAlias = np.ndarray | list | tuple +OriginLike: TypeAlias = VectorLike | pd.Series | str | None + + +def _default_axes(data: 'SarracenDataFrame', + x: str | None, + y: str | None) -> tuple[str, str]: """ Utility function to determine the x & y columns to use for rendering. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to render. - x, y: str + x, y : str The x and y directional column labels passed to the render function. Returns ------- - x, y: str + x, y : str The directional column labels to use for rendering. Defaults to the x-column detected in `data` """ @@ -52,28 +60,26 @@ def _default_axes(data: 'SarracenDataFrame', # noqa: F821 return x, y -def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 +def _rotate_data(data: 'SarracenDataFrame', x_data: np.ndarray, y_data: np.ndarray, z_data: np.ndarray, - rotation: Union[np.ndarray, list, Rotation, None], - rot_origin: Union[np.ndarray, list, pd.Series, - str, None]) -> Tuple[np.ndarray, - np.ndarray, - np.ndarray]: + rotation: VectorLike | Rotation | None, + rot_origin: OriginLike) -> tuple[np.ndarray, np.ndarray, + np.ndarray]: """ Rotate vector data in a particle dataset. Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame The particle dataset to interpolate over. - x_data, y_data, z_data: ndarray + x_data, y_data, z_data : ndarray The directional vector data. - rotation: array_like or SciPy Rotation + rotation : array_like or SciPy Rotation The rotation to apply to the vector data. If defined as an array, the order of rotations is [z, y, x] in degrees - rot_origin: array_like or ['com', 'midpoint'] + rot_origin : array_like or ['com', 'midpoint'] Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the @@ -82,7 +88,7 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 Returns ------- - x_data, y_data, z_data: ndarray + x_data, y_data, z_data : ndarray The rotated x, y, and z directional data. """ vectors = np.array([x_data, y_data, z_data]).transpose() @@ -101,7 +107,7 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 rot_origin_arr = data.centre_of_mass() elif rot_origin == 'midpoint': rot_origin_arr = (vectors.min(0) + vectors.max(0)) / 2 - elif not isinstance(rot_origin, (list, pd.Series, np.ndarray)): + elif not isinstance(rot_origin, (list, tuple, pd.Series, np.ndarray)): raise ValueError("rot_origin should be an [x, y, z] point or " "'com' or 'midpoint'") elif len(rot_origin) != 3: @@ -115,12 +121,12 @@ def _rotate_data(data: 'SarracenDataFrame', # noqa: F821 return vectors[:, 0], vectors[:, 1], vectors[:, 2] -def _default_bounding_box(data: 'SarracenDataFrame', # noqa: F821 +def _default_bounding_box(data: 'SarracenDataFrame', x: str, y: str, - xlim: Union[Tuple[float, float], None], - ylim: Union[Tuple[float, float], None], - z_slice: Union[float, None] = None) -> np.ndarray: + xlim: Bounds | None, + ylim: Bounds | None, + z_slice: float | None = None) -> np.ndarray: # boundaries of the plot default to the max & min values of the data. x_min = xlim[0] if xlim is not None and xlim[0] is not None else None y_min = ylim[0] if ylim is not None and ylim[0] is not None else None @@ -141,33 +147,28 @@ def _default_bounding_box(data: 'SarracenDataFrame', # noqa: F821 return np.array(corners).transpose() -def _default_bounds(data: 'SarracenDataFrame', # noqa: F821 - x_data: np.ndarray, - y_data: np.ndarray, - xlim: Union[Tuple[float, float], None], - ylim: Union[Tuple[float, float], None], - ) -> Tuple[Tuple[float, float], - Tuple[float, float]]: +def _default_bounds(x_data: np.ndarray | pd.Series, + y_data: np.ndarray | pd.Series, + xlim: Bounds | None, + ylim: Bounds | None) -> tuple[tuple[float, float], + tuple[float, float]]: """ Utility function to determine the 2-dimensional boundaries to use in 2D rendering. Parameters ---------- - data: SarracenDataFrame - The particle dataset to render. - x, y: str - The directional column labels that will be used for rendering. - xlim, ylim: tuple of float + x_data, y_data : ndarray or Series + The x and y particle coordinate data. + xlim, ylim : tuple of float or None, optional The minimum and maximum values passed to the render function, in particle data space. Returns ------- - xlim, ylim: tuple of float + xlim, ylim : tuple of float The minimum and maximum values to use for rendering, in particle data - space. Defaults to the maximum and minimum values of `x` and `y`, - snapped to the nearest integer. + space. Defaults to the maximum and minimum values of `x` and `y`. """ # boundaries of the plot default to the max & min values of the data. x_min = xlim[0] if xlim is not None and xlim[0] is not None else None @@ -183,27 +184,27 @@ def _default_bounds(data: 'SarracenDataFrame', # noqa: F821 return (x_min, x_max), (y_min, y_max) -def _set_pixels(x_pixels: Union[int, None], - y_pixels: Union[int, None], - xlim: Tuple[float, float], - ylim: Tuple[float, float], - default: int) -> Tuple[int, int]: +def _set_pixels(x_pixels: int | None, + y_pixels: int | None, + xlim: tuple[float, float], + ylim: tuple[float, float], + default: int) -> tuple[int, int]: """ Utility function to determine the number of pixels to interpolate over in 2D interpolation. Parameters ---------- - x_pixels, y_pixels: int + x_pixels, y_pixels : int The number of pixels in the x & y directions passed to the interpolation function. - xlim, ylim: tuple of float + xlim, ylim : tuple of float The minimum and maximum values to use in interpolation, in particle data space. Returns ------- - x_pixels, y_pixels: int + x_pixels, y_pixels : int The number of pixels in the x & y directions to use in 2D interpolation. """ @@ -223,33 +224,33 @@ def _set_pixels(x_pixels: Union[int, None], return x_pixels, y_pixels -def render(data: 'SarracenDataFrame', # noqa: F821 +def render(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[float, None] = None, - kernel: Union[BaseKernel, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - cmap: Union[str, Colormap] = 'gist_heat', + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: float | None = None, + kernel: BaseKernel | None = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + cmap: str | Colormap = 'gist_heat', cbar: bool = True, cbar_kws: dict = {}, - cbar_ax: Union[Axes, None] = None, - ax: Union[Axes, None] = None, + cbar_ax: Axes | None = None, + ax: Axes | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, log_scale: bool = False, symlog_scale: bool = False, - dens_weight: Union[bool, None] = None, + dens_weight: bool | None = None, normalize: bool = True, hmin: bool = False, - corotation: Union[np.ndarray, list, None] = None, + corotation: VectorLike | None = None, **kwargs: Any) -> Axes: """ Render a scalar SPH target variable to a grid plot. @@ -258,75 +259,78 @@ def render(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target variable. - x, y, z: str, optional + x, y, z : str, optional Column labels of the x, y & z directional axes. Defaults to the columns detected in `data`. - xsec: float, optional. + xsec : float, optional For a 3D dataset, the z value to take a cross-section at. If none, column interpolation is performed. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of pixels present in the final image. - xlim, ylim: tuple of float, optional - The starting and ending corners of the final 2D image. - vmin, vmax: float, optional + xlim, ylim : tuple of float or None, optional + The starting and ending corners of the 2D image. A tuple gives the + upper and lower bounds. Either value of the tuple may be `None`, which + defaults to the minimum or maximum value of `x` and `y`, respectively. + If omitted, the limit defaults to the minimum and maximum values. + vmin, vmax : float, optional Lower and upper limits of the range of values for the colour bar. - cmap: str or Colormap, optional + cmap : str or Colormap, optional The color map to use when plotting a 2D image. - cbar: bool, optional + cbar : bool, optional True if a colorbar should be drawn. - cbar_kws: dict, optional + cbar_kws : dict, optional Keyword arguments to pass to matplotlib.figure.Figure.colorbar(). - cbar_ax: Axes, optional + cbar_ax : Axes, optional Axes to draw the colorbar in, if not provided then space will be taken from the main Axes. - ax: Axes, optional + ax : Axes, optional The main axes in which to draw the rendered image. - exact: bool, optional + exact : bool, optional Whether to use exact interpolation of the data. For cross-sections this is ignored. Defaults to False. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - integral_samples: int, optional + integral_samples : int, optional If using column interpolation, the number of sample points to take when approximating the 2D column kernel. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. Only applies to 3D datasets. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - log_scale: bool, optional + log_scale : bool, optional Whether to use a logarithmic scale for color coding. - symlog_scale: bool, optional + symlog_scale : bool, optional Whether to use a symmetrical logarithmic scale for color coding (i.e., allows positive and negative values). Optionally add "linthresh" and "linscale" to kwargs to set the linear region and the scaling of linear values, respectively (defaults to 1e-9 and 1, respectively). Only works if log_scale == True. - dens_weight: bool, optional + dens_weight : bool, optional If True, will plot the target multiplied by the density. Defaults to True for column-integrated views and False for everything else. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to True. - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). - corotation: list, optional - Moves particles to the co-rotating frame of two location. corotation + corotation : array_like, optional + Moves particles to the co-rotating frame of two location. Corotation contains two lists which correspond to the two x, y, z coordinates. - kwargs: other keyword arguments + kwargs : other keyword arguments Keyword arguments to pass to ax.imshow. Returns @@ -352,8 +356,8 @@ def render(data: 'SarracenDataFrame', # noqa: F821 A_{pixel} = \\sum_b \\frac{m_b}{\\rho_b} A_b W_{ab}(h_b) - where :math:`m` is the mass, :math:`\\rho` is the density, and :math:`W` is - the smoothing kernel with smoothing length, :math:`h`. + where :math:`m` is the mass, :math:`\\rho` is the density, and :math:`W` + is the smoothing kernel with smoothing length, :math:`h`. Normalized interpolation divides the above summation by an interpolation of a constant scalar field equal to 1: @@ -428,7 +432,7 @@ def render(data: 'SarracenDataFrame', # noqa: F821 corners = _default_bounding_box(data, x, y, xlim, ylim, xsec) rotated_corners = _rotate_data(data, corners[0], corners[1], corners[2], rotation, rot_origin) - xlim, ylim = _default_bounds(data, rotated_corners[0], rotated_corners[1], + xlim, ylim = _default_bounds(rotated_corners[0], rotated_corners[1], xlim, ylim) kwargs.setdefault("origin", 'lower') @@ -473,18 +477,18 @@ def render(data: 'SarracenDataFrame', # noqa: F821 return ax -def lineplot(data: 'SarracenDataFrame', # noqa: F821 +def lineplot(data: 'SarracenDataFrame', target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, pixels: int = 512, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - zlim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, - backend: Union[str, None] = None, + xlim: Bounds | float | int | None = None, + ylim: Bounds | float | int | None = None, + zlim: Bounds | float | int | None = None, + ax: Axes | None = None, + backend: str | None = None, log_scale: bool = False, dens_weight: bool = False, normalize: bool = True, @@ -497,37 +501,42 @@ def lineplot(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str + target : str Column label of the target variable. - x, y, z: str, optional + x, y, z : str, optional Column labels of the x, y & z directional axes. Defaults to the columns detected in `data`. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - pixels: int, optional + pixels : int, optional Number of samples taken across the x axis in the final plot. - xlim, ylim, zlim: tuple of float, optional - Coordinates of the two points that make up the cross-sectional line. - ax: Axes, optional + xlim, ylim, zlim : int, float, or tuple of float or None, optional + Coordinate bounds of the cross-sectional line. A scalar value is + treated as both bounds for that axis, fixing the line along that + constant coordinate value. A tuple gives the upper and lower bounds. + Either value of the tuple may be `None`, which defaults to the minimum + or maximum value of `x`, `y`, and `z`, respectively. If omitted, the + limit defaults to the minimum and maximum values. + ax : Axes, optional The main axes in which to draw the final plot. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - log_scale: bool, optional + log_scale : bool, optional Whether to use a logarithmic scale for color coding. - dens_weight: bool, optional + dens_weight : bool, optional If True, will plot the target mutliplied by the density. Defaults to False. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). - kwargs: other keyword arguments + kwargs : other keyword arguments Keyword arguments to pass to sns.lineplot. Returns @@ -538,8 +547,8 @@ def lineplot(data: 'SarracenDataFrame', # noqa: F821 Raises ------- ValueError - If `x_pixels` or `y_pixels` are less than or equal to zero, or if the - specified `x` and `y` minimum and maximums result in an invalid region. + If `pixels` are less than or equal to zero, or if the specified `x` + and `y` minimum and maximum values result in an invalid region. KeyError If `target`, `x`, `y`, mass, density, or smoothing length columns do not exist in `data`. @@ -559,7 +568,7 @@ def lineplot(data: 'SarracenDataFrame', # noqa: F821 ylim = ylim, ylim x, y = _default_axes(data, x, y) - xlim, ylim = _default_bounds(data, data.loc[:, x], data.loc[:, y], + xlim, ylim = _default_bounds(data.loc[:, x], data.loc[:, y], xlim, ylim) if data.get_dim() == 2: @@ -578,6 +587,9 @@ def lineplot(data: 'SarracenDataFrame', # noqa: F821 z2 = zmin if zlim is None or zlim[1] is None else zlim[1] zlim = z2, z1 + if zlim is None or zlim[0] is None or zlim[1] is None: + raise ValueError("z-limits must be specified.") + upper_lim = np.sqrt((xlim[1] - xlim[0])**2 + (ylim[1] - ylim[0])**2 + (zlim[1] - zlim[0])**2) @@ -602,24 +614,24 @@ def lineplot(data: 'SarracenDataFrame', # noqa: F821 return ax -def streamlines(data: 'SarracenDataFrame', # noqa: F821 - target: Union[Tuple[str, str], Tuple[str, str, str]], - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[float, None] = None, - kernel: Union[BaseKernel, None] = None, +def streamlines(data: 'SarracenDataFrame', + target: tuple[str, str] | tuple[str, str, str], + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: float | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + ax: Axes | None = None, exact: bool = False, - backend: Union[str, None] = None, - dens_weight: Union[bool, None] = None, + backend: str | None = None, + dens_weight: bool | None = None, normalize: bool = True, hmin: bool = False, **kwargs: Any) -> Axes: @@ -634,58 +646,61 @@ def streamlines(data: 'SarracenDataFrame', # noqa: F821 Parameters ---------- - data: SarracenDataFrame + data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str tuple of shape (2) or (3). + target : str tuple of shape (2) or (3). Column label of the target vector. Shape must match the # of dimensions in `data`. - x, y, z: str, optional + x, y, z : str, optional Column label of the x, y & z directional axes. Defaults to the columns detected in `data`. - xsec: float, optional + xsec : float, optional The z to take a cross-section at. If none, column interpolation is performed. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - integral_samples: int, optional + integral_samples : int, optional If using column interpolation, the number of sample points to take when approximating the 2D column kernel. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels: int, optional + x_pixels, y_pixels : int, optional Number of interpolation samples to pass to ax.streamlines(). Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: float, optional - The minimum and maximum values to use in interpolation, in particle - data space. Defaults to the minimum and maximum values of `x` and `y`. - ax: Axes, optional + xlim, ylim : tuple of float or None, optional + The minimum and maximum values to use in interpolation along each axis. + A tuple gives the upper and lower bounds. Either value of the tuple may + be `None`, which defaults to the minimum or maximum value of `x` and + `y`, respectively. If omitted, the limit defaults to the minimum and + maximum values. + ax : Axes, optional The main axes in which to draw the rendered image. - exact: bool, optional + exact : bool, optional Whether to use exact interpolation of the data. For cross-sections this is ignored. Defaults to False. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, will plot the target mutliplied by the density. Defaults to True for column-integrated views and False for everything else. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). - kwargs: other keyword arguments + kwargs : other keyword arguments Keyword arguments to pass to ax.streamlines() Returns @@ -743,7 +758,7 @@ def streamlines(data: 'SarracenDataFrame', # noqa: F821 corners = _default_bounding_box(data, x, y, xlim, ylim, xsec) rotated_corners = _rotate_data(data, corners[0], corners[1], corners[2], rotation, rot_origin) - xlim, ylim = _default_bounds(data, rotated_corners[0], rotated_corners[1], + xlim, ylim = _default_bounds(rotated_corners[0], rotated_corners[1], xlim, ylim) kwargs.setdefault("color", 'black') @@ -763,26 +778,26 @@ def streamlines(data: 'SarracenDataFrame', # noqa: F821 return ax -def arrowplot(data: 'SarracenDataFrame', # noqa: F821 - target: Union[Tuple[str, str], Tuple[str, str, str]], - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[float, None] = None, - kernel: Union[BaseKernel, None] = None, +def arrowplot(data: 'SarracenDataFrame', + target: tuple[str, str] | tuple[str, str, str], + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: float | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_arrows: Union[int, None] = None, - y_arrows: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_arrows: int | None = None, + y_arrows: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + ax: Axes | None = None, qkey: bool = True, - qkey_kws: Union[dict, None] = None, + qkey_kws: dict | None = None, exact: bool = False, - backend: Union[str, None] = None, - dens_weight: Union[bool, None] = None, + backend: str | None = None, + dens_weight: bool | None = None, normalize: bool = True, hmin: bool = False, **kwargs: Any) -> Axes: @@ -799,60 +814,63 @@ def arrowplot(data: 'SarracenDataFrame', # noqa: F821 ---------- data : SarracenDataFrame Particle data, in a SarracenDataFrame. - target: str tuple of shape (2) or (3). + target : str tuple of shape (2) or (3). Column label of the target vector. Shape must match the # of dimensions in `data`. - x, y, z: str, optional + x, y, z : str, optional Column label of the x, y & z directional axes. Defaults to the columns detected in `data`. - xsec: float, optional + xsec : float, optional The z to take a cross-section at. If none, column interpolation is performed. - kernel: BaseKernel, optional + kernel : BaseKernel, optional Kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - integral_samples: int, optional + integral_samples : int, optional If using column interpolation, the number of sample points to take when approximating the 2D column kernel. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_arrows, y_arrows: int, optional + x_arrows, y_arrows : int, optional Number of arrows in the output image in the x & y directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim: tuple of float, optional - The minimum and maximum values to use in interpolation, in particle - data space. Defaults to the minimum and maximum values of `x` and `y`. - ax: Axes, optional + xlim, ylim : tuple of float or None, optional + The minimum and maximum values to use in interpolation along each axis. + A tuple gives the upper and lower bounds. Either value of the tuple may + be `None`, which defaults to the minimum or maximum value of `x` and + `y`, respectively. If omitted, the limit defaults to the minimum and + maximum values. + ax : Axes, optional The main axes in which to draw the rendered image. - qkey: bool, optional + qkey : bool, optional Whether to include a quiver key on the final plot. - qkey_kws: dict, optional + qkey_kws : dict, optional Keywords to pass through to ax.quiver. - exact: bool, optional + exact : bool, optional Whether to use exact interpolation of the data. For cross-sections this is ignored. Defaults to False. - backend: ['cpu', 'gpu'], optional + backend : ['cpu', 'gpu'], optional The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool, optional + dens_weight : bool, optional If True, will plot the target mutliplied by the density. Defaults to True for column-integrated views and False for everything else. - normalize: bool, optional + normalize : bool, optional If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool, optional + hmin : bool, optional If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future verison). - kwargs: other keyword arguments + kwargs : other keyword arguments Keyword arguments to pass to ax.quiver() Returns @@ -875,7 +893,7 @@ def arrowplot(data: 'SarracenDataFrame', # noqa: F821 corners = _default_bounding_box(data, x, y, xlim, ylim, xsec) rotated_corners = _rotate_data(data, corners[0], corners[1], corners[2], rotation, rot_origin) - xlim, ylim = _default_bounds(data, rotated_corners[0], rotated_corners[1], + xlim, ylim = _default_bounds(rotated_corners[0], rotated_corners[1], xlim, ylim) x_arrows, y_arrows = _set_pixels(x_arrows, y_arrows, xlim, ylim, 20) diff --git a/sarracen/sarracen_dataframe.py b/sarracen/sarracen_dataframe.py index 4b242d9..435e83a 100644 --- a/sarracen/sarracen_dataframe.py +++ b/sarracen/sarracen_dataframe.py @@ -1,4 +1,5 @@ -from typing import Any, Type, Union, Callable, Tuple, Optional, Dict, List +from collections.abc import Callable +from typing import Any, TypeAlias from matplotlib.axes import Axes from matplotlib.colors import Colormap @@ -13,6 +14,11 @@ from .kernels import CubicSplineKernel, BaseKernel +Bounds: TypeAlias = tuple[float | None, float | None] +VectorLike: TypeAlias = np.ndarray | list | tuple +OriginLike: TypeAlias = VectorLike | pd.Series | str | None + + def _copy_doc(copy_func: Callable) -> Callable: """Copy documentation from another function to this function.""" def wrapper(func: Callable) -> Callable: @@ -124,7 +130,7 @@ def __init__(self, data=None, params=None, *args, **kwargs): self._backend = 'gpu' if cuda.is_available() else 'cpu' @property - def _constructor(self) -> Type: + def _constructor(self) -> type: return SarracenDataFrame def _identify_special_columns(self) -> None: @@ -317,28 +323,28 @@ def centre_of_mass(self) -> list: @_copy_doc(render) def render(self, target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[float, None] = None, - kernel: Union[BaseKernel, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - cmap: Union[str, Colormap] = 'gist_heat', + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: float | None = None, + kernel: BaseKernel | None = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + cmap: str | Colormap = 'gist_heat', cbar: bool = True, cbar_kws: dict = {}, - cbar_ax: Union[Axes, None] = None, - ax: Union[Axes, None] = None, + cbar_ax: Axes | None = None, + ax: Axes | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, log_scale: bool = False, symlog_scale: bool = False, - dens_weight: Union[bool, None] = None, + dens_weight: bool | None = None, normalize: bool = False, hmin: bool = False, **kwargs: Any) -> Axes: @@ -351,16 +357,16 @@ def render(self, @_copy_doc(lineplot) def lineplot(self, target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, pixels: int = 512, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - zlim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, - backend: Union[str, None] = None, + xlim: Bounds | float | int | None = None, + ylim: Bounds | float | int | None = None, + zlim: Bounds | float | int | None = None, + ax: Axes | None = None, + backend: str | None = None, log_scale: bool = False, dens_weight: bool = False, normalize: bool = False, @@ -372,22 +378,22 @@ def lineplot(self, @_copy_doc(streamlines) def streamlines(self, - target: Union[Tuple[str, str], Tuple[str, str, str]], - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[int, None] = None, - kernel: Union[BaseKernel, None] = None, + target: tuple[str, str] | tuple[str, str, str], + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: int | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + ax: Axes | None = None, exact: bool = False, - backend: Union[str, None] = None, + backend: str | None = None, dens_weight: bool = False, normalize: bool = False, hmin: bool = False, @@ -399,25 +405,25 @@ def streamlines(self, @_copy_doc(arrowplot) def arrowplot(self, - target: Union[Tuple[str, str], Tuple[str, str, str]], - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - xsec: Union[int, None] = None, - kernel: Union[BaseKernel, None] = None, + target: tuple[str, str] | tuple[str, str, str], + x: str | None = None, + y: str | None = None, + z: str | None = None, + xsec: int | None = None, + kernel: BaseKernel | None = None, integral_samples: int = 1000, - rotation: Union[np.ndarray, list, Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_arrows: Union[int, None] = None, - y_arrows: Union[int, None] = None, - xlim: Union[Tuple[float, float], None] = None, - ylim: Union[Tuple[float, float], None] = None, - ax: Union[Axes, None] = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_arrows: int | None = None, + y_arrows: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + ax: Axes | None = None, qkey: bool = True, - qkey_kws: Union[dict, None] = None, + qkey_kws: dict | None = None, exact: bool = False, - backend: Union[str, None] = None, - dens_weight: Union[bool, None] = None, + backend: str | None = None, + dens_weight: bool | None = None, normalize: bool = False, hmin: bool = False, **kwargs: Any) -> Axes: @@ -428,21 +434,18 @@ def arrowplot(self, def sph_interpolate(self, target: str, - x: Union[str, None] = None, - y: Union[str, None] = None, - z: Union[str, None] = None, - kernel: Union[BaseKernel, None] = None, - rotation: Union[np.ndarray, list, - Rotation, None] = None, - rot_origin: Union[np.ndarray, list, str, None] = None, - x_pixels: Union[int, None] = None, - y_pixels: Union[int, None] = None, - z_pixels: Union[int, None] = None, - xlim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - ylim: Optional[Tuple[Optional[float], - Optional[float]]] = None, - zlim: Union[Tuple[float, float], None] = None, + x: str | None = None, + y: str | None = None, + z: str | None = None, + kernel: BaseKernel | None = None, + rotation: VectorLike | Rotation | None = None, + rot_origin: OriginLike = None, + x_pixels: int | None = None, + y_pixels: int | None = None, + z_pixels: int | None = None, + xlim: Bounds | None = None, + ylim: Bounds | None = None, + zlim: Bounds | None = None, exact: bool = False, backend: str = 'cpu', dens_weight: bool = False, @@ -454,45 +457,45 @@ def sph_interpolate(self, Parameters ---------- - target: str + target : str The column label of the target data. - x, y, z: str + x, y, z : str The column labels of the directional data to interpolate over. Defaults to the x, y, and z columns detected in `data`. - kernel: BaseKernel + kernel : BaseKernel The kernel to use for smoothing the target data. Defaults to the kernel specified in `data`. - rotation: array_like or SciPy Rotation, optional + rotation : array_like or SciPy Rotation, optional The rotation to apply to the data before interpolation. If defined as an array, the order of rotations is [z, y, x] in degrees. Only applies to 3D datasets. - rot_origin: array_like or ['com', 'midpoint'], optional + rot_origin : array_like or ['com', 'midpoint'], optional Point of rotation of the data. Only applies to 3D datasets. If array_like, then the [x, y, z] coordinates specify the point around which the data is rotated. If 'com', then data is rotated around the centre of mass. If 'midpoint', then data is rotated around the midpoint, that is, min + max / 2. Defaults to the midpoint. - x_pixels, y_pixels, z_pixels: int, optional + x_pixels, y_pixels, z_pixels : int, optional Number of pixels in the output image in the x, y & z directions. Default values are chosen to keep a consistent aspect ratio. - xlim, ylim, zlim: tuple of float, optional + xlim, ylim, zlim : tuple of float or None, optional The minimum and maximum values to use in interpolation, in particle data space. Defaults to the minimum and maximum values of `x`, `y` and `z`. - exact: bool + exact : bool Whether to use exact interpolation of the data. Only applies to 2D datasets. - backend: ['cpu', 'gpu'] + backend : ['cpu', 'gpu'] The computation backend to use when interpolating this data. Defaults to 'gpu' if CUDA is enabled, otherwise 'cpu' is used. A manually specified backend in `data` will override the default. - dens_weight: bool + dens_weight : bool If True, the target will be multiplied by density. Defaults to False. - normalize: bool + normalize : bool If True, will normalize the interpolation. Defaults to False (this may change in future versions). - hmin: bool + hmin : bool If True, a minimum smoothing length of 0.5 * pixel size will be imposed. This ensures each particle contributes to at least one grid cell / pixel. Defaults to False (this may change in a future @@ -535,10 +538,6 @@ def sph_interpolate(self, (165, 526, 512) """ if self.get_dim() == 2: - if xlim is None: - xlim = (None, None) - if ylim is None: - ylim = (None, None) return interpolate_2d(self, target, x, y, kernel, x_pixels, y_pixels, xlim, ylim, exact, backend, dens_weight, normalize, hmin) @@ -550,7 +549,7 @@ def sph_interpolate(self, raise ValueError('Invalid number of dimensions.') @property - def params(self) -> Dict[str, Any]: + def params(self) -> dict[str, Any]: """ dict: Miscellaneous dataset-level parameters. @@ -562,14 +561,14 @@ def params(self) -> Dict[str, Any]: return self._params @params.setter - def params(self, new_params: Union[Dict[str, Any], None]) -> None: + def params(self, new_params: dict[str, Any] | None) -> None: if new_params is not None and not isinstance(new_params, dict): raise TypeError("Parameters not a dictionary") self._params = dict(new_params or {}) @property def units(self) -> Series: - """Series: Units for each column of this dataset.""" + """Series : Units for each column of this dataset.""" return self._units @units.setter @@ -587,7 +586,7 @@ def xcol(self) -> str: return self._xcol @xcol.setter - def xcol(self, new_col: Union[str, None]) -> None: + def xcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._xcol = new_col @@ -602,7 +601,7 @@ def ycol(self) -> str: return self._ycol @ycol.setter - def ycol(self, new_col: Union[str, None]) -> None: + def ycol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._ycol = new_col @@ -617,7 +616,7 @@ def zcol(self) -> str: return self._zcol @zcol.setter - def zcol(self, new_col: Union[str, None]) -> None: + def zcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._zcol = new_col @@ -632,7 +631,7 @@ def hcol(self) -> str: return self._hcol @hcol.setter - def hcol(self, new_col: Union[str, None]) -> None: + def hcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._hcol = new_col @@ -647,7 +646,7 @@ def mcol(self) -> str: return self._mcol @mcol.setter - def mcol(self, new_col: Union[str, None]) -> None: + def mcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._mcol = new_col @@ -662,7 +661,7 @@ def rhocol(self) -> str: return self._rhocol @rhocol.setter - def rhocol(self, new_col: Union[str, None]) -> None: + def rhocol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._rhocol = new_col @@ -678,7 +677,7 @@ def vxcol(self) -> str: return self._vxcol @vxcol.setter - def vxcol(self, new_col: Union[str, None]) -> None: + def vxcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._vxcol = new_col @@ -694,7 +693,7 @@ def vycol(self) -> str: return self._vycol @vycol.setter - def vycol(self, new_col: Union[str, None]) -> None: + def vycol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._vycol = new_col @@ -710,16 +709,16 @@ def vzcol(self) -> str: return self._vzcol @vzcol.setter - def vzcol(self, new_col: Union[str, None]) -> None: + def vzcol(self, new_col: str | None) -> None: if new_col in self or new_col is None: self._vzcol = new_col @property - def dustfracscol(self) -> List[str]: + def dustfracscol(self) -> list[str]: return self._dustfracscol @dustfracscol.setter - def dustfracscol(self, new_col: List[str]) -> None: + def dustfracscol(self, new_col: list[str]) -> None: if new_col in self or new_col is None: self._dustfracscol = new_col diff --git a/sarracen/tests/disc/test_surface_density.py b/sarracen/tests/disc/test_surface_density.py index e4507de..97d6d45 100644 --- a/sarracen/tests/disc/test_surface_density.py +++ b/sarracen/tests/disc/test_surface_density.py @@ -124,9 +124,20 @@ def test_one_fluid_single_grain() -> None: 'vx': vx, 'vy': vy, 'vz': vz, 'dustfrac': dustfrac, 'mass': mass}) - sigma_g_in, sigma_d_in = surface_density(sdf, r_in=0, r_out=0.5, bins=50) - sigma_g_out, sigma_d_out = surface_density(sdf, r_in=0.5, r_out=1, bins=50) - sigma_g_all, sigma_d_all = surface_density(sdf, r_in=0, r_out=1, bins=100) + sigmas = surface_density(sdf, r_in=0, r_out=0.5, bins=50) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 2 + sigma_g_in, sigma_d_in = sigmas + + sigmas = surface_density(sdf, r_in=0.5, r_out=1, bins=50) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 2 + sigma_g_out, sigma_d_out = sigmas + + sigmas = surface_density(sdf, r_in=0, r_out=1, bins=100) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 2 + sigma_g_all, sigma_d_all = sigmas assert_allclose(sigma_g_in, sigma_g_all[:50], atol=1e-15, rtol=0.0) assert_allclose(sigma_g_out, sigma_g_all[50:], atol=1e-15, rtol=0.0) @@ -163,12 +174,18 @@ def test_one_fluid_multi_grain() -> None: sdf.params['ndustsmall'] = 3 sigmas = surface_density(sdf, r_in=0, r_out=0.5, bins=50) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 3 sigma_g_in, sigma_d_tot_in, sigma_d_in = sigmas sigmas = surface_density(sdf, r_in=0.5, r_out=1, bins=50) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 3 sigma_g_out, sigma_d_tot_out, sigma_d_out = sigmas sigmas = surface_density(sdf, r_in=0, r_out=1, bins=100) + assert isinstance(sigmas, tuple) + assert len(sigmas) == 3 sigma_g_all, sigma_d_tot_all, sigma_d_all = sigmas assert_allclose(sigma_g_in, sigma_g_all[:50], atol=1e-15, rtol=0) diff --git a/sarracen/tests/interpolate/test_interpolate.py b/sarracen/tests/interpolate/test_interpolate.py index 51532b3..333a0fb 100644 --- a/sarracen/tests/interpolate/test_interpolate.py +++ b/sarracen/tests/interpolate/test_interpolate.py @@ -1,7 +1,8 @@ """ pytest unit tests for interpolate.py functions. """ -from typing import Any, Callable, Dict, List +from collections.abc import Callable +from typing import Any import pandas as pd import numpy as np @@ -21,14 +22,14 @@ if cuda.is_available(): backends.append('gpu') -funcs2d: List[Callable] = [interpolate_2d, interpolate_2d_line] -funcs2dvec: List[Callable] = [interpolate_2d_vec] -funcs3d: List[Callable] = [interpolate_3d_line, interpolate_3d_proj, +funcs2d: list[Callable] = [interpolate_2d, interpolate_2d_line] +funcs2dvec: list[Callable] = [interpolate_2d_vec] +funcs3d: list[Callable] = [interpolate_3d_line, interpolate_3d_proj, interpolate_3d_cross, interpolate_3d_grid] -funcs3dvec: List[Callable] = [interpolate_3d_vec, interpolate_3d_cross_vec] +funcs3dvec: list[Callable] = [interpolate_3d_vec, interpolate_3d_cross_vec] -funcscolumn: List[Callable] = [interpolate_3d_proj, interpolate_3d_vec] -funcsline: List[Callable] = [interpolate_2d_line, interpolate_3d_line] +funcscolumn: list[Callable] = [interpolate_3d_proj, interpolate_3d_vec] +funcsline: list[Callable] = [interpolate_2d_line, interpolate_3d_line] funcs = funcs2d + funcs2dvec + funcs3d + funcs3dvec @@ -624,7 +625,7 @@ def test_default_kernel(backend: str, use_default_kernel: bool) -> None: sdf_2.backend = backend sdf_3.backend = backend - kwargs: Dict[str, Any] = {'normalize': False} + kwargs: dict[str, Any] = {'normalize': False} if use_default_kernel: kernel: BaseKernel = QuarticSplineKernel() @@ -776,7 +777,7 @@ def test_pixel_arguments() -> None: assert img.shape == (default_pixels, default_pixels, default_pixels) # Non-vector functions - functions_list: List[Callable] = [interpolate_2d, interpolate_3d_proj, + functions_list: list[Callable] = [interpolate_2d, interpolate_3d_proj, interpolate_3d_cross] for func in functions_list: for axes in [('x', 'y'), @@ -1200,7 +1201,7 @@ def test_required_columns(backend: str, func: Callable, column: str) -> None: sdf = SarracenDataFrame(data, params=dict()) sdf.backend = backend - kwargs: Dict[str, Any] = dict() + kwargs: dict[str, Any] = dict() if func in funcs2d + funcs3d: kwargs['target'] = 'A' elif func in funcs2dvec: @@ -1289,7 +1290,7 @@ def test_density_weighted(backend: str, if not dens_weight: weight = weight / sdf['rho'][0] - kwargs: Dict[str, Any] = {'xlim': (-1, 1), 'ylim': (-1, 1), + kwargs: dict[str, Any] = {'xlim': (-1, 1), 'ylim': (-1, 1), 'dens_weight': dens_weight, 'normalize': False, 'hmin': False} @@ -1348,7 +1349,7 @@ def test_normalize_interpolation(backend: str, normalize: bool) -> None: norm3d = weight / sdf_2['h'][0] * kernel.w(0, 3) norm3d_column = weight * kernel.get_column_kernel()[0] - kwargs: Dict[str, Any] = {'xlim': (-1, 1), 'ylim': (-1, 1), + kwargs: dict[str, Any] = {'xlim': (-1, 1), 'ylim': (-1, 1), 'dens_weight': False, 'normalize': normalize} img = interpolate_2d(sdf_2, 'A', @@ -1478,7 +1479,7 @@ def test_minimum_smoothing_length_3d(backend: str) -> None: sdf_a.backend = backend sdf_b.backend = backend - functions_list: List[Callable] = [interpolate_3d_cross, + functions_list: list[Callable] = [interpolate_3d_cross, interpolate_3d_proj, interpolate_3d_grid] for interpolate in functions_list: diff --git a/sarracen/tests/interpolate/test_rotation.py b/sarracen/tests/interpolate/test_rotation.py index bd95bc8..b2a09cd 100644 --- a/sarracen/tests/interpolate/test_rotation.py +++ b/sarracen/tests/interpolate/test_rotation.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Tuple +from collections.abc import Callable import numpy as np from numba import cuda @@ -16,18 +16,18 @@ backends.append('gpu') -def rotate(target: Tuple[float, float, float], +def rotate(target: tuple[float, float, float], rot_z: float, rot_y: float, - rot_x: float) -> Tuple[float, float, float]: + rot_x: float) -> tuple[float, float, float]: """ Perform a rotation of a target vector in three dimensions. A helper function for test_nonstandard_rotation() Parameters ---------- - target: float tuple of shape (3) - rot_z, rot_y, rot_x: Rotation around each axis (in degrees) + target : float tuple of shape (3) + rot_z, rot_y, rot_x : Rotation around each axis (in degrees) Returns ------- @@ -342,7 +342,7 @@ def test_rotation_stability(backend: str) -> None: real = -1 + (np.arange(0, 50) + 0.5) * (1 / 25) pixel_x, pixel_y = 12, 30 - functions_list: List[Callable] = [interpolate_3d_proj, + functions_list: list[Callable] = [interpolate_3d_proj, interpolate_3d_cross] for func in functions_list: img = func(sdf, 'A', @@ -463,7 +463,7 @@ def test_axes_rotation_equivalency(backend: str) -> None: for i_x in range(4): rot_x, rot_y, rot_z = i_x * 90, i_y * 90, i_z * 90 - functions_list: List[Callable] = [interpolate_3d_proj, + functions_list: list[Callable] = [interpolate_3d_proj, interpolate_3d_cross] for func in functions_list: img1 = func(sdf, 'A', @@ -518,7 +518,7 @@ def test_com_rotation() -> None: sdf_zero = SarracenDataFrame({'x': x, 'y': y, 'z': z, 'h': h, 'val': val}, params={'mass': mass, 'hfact': 1.2}) - functions_list: List[Callable] = [interpolate_3d_proj, + functions_list: list[Callable] = [interpolate_3d_proj, interpolate_3d_cross] for func in functions_list: img1 = func(sdf_com, 'val', diff --git a/sarracen/tests/kernels/test_kernels.py b/sarracen/tests/kernels/test_kernels.py index 7d4fd94..74d2fcc 100644 --- a/sarracen/tests/kernels/test_kernels.py +++ b/sarracen/tests/kernels/test_kernels.py @@ -1,5 +1,6 @@ """pytest unit tests for kernel functionality.""" -from typing import Callable +from collections.abc import Callable + from pytest import approx, mark from scipy.integrate import quad, dblquad, tplquad import numpy as np diff --git a/sarracen/tests/readers/test_read_phantom.py b/sarracen/tests/readers/test_read_phantom.py index 297ef30..88e4e87 100644 --- a/sarracen/tests/readers/test_read_phantom.py +++ b/sarracen/tests/readers/test_read_phantom.py @@ -1,4 +1,3 @@ -from typing import Type, Union, List, Tuple, Dict import pandas as pd import numpy as np @@ -10,8 +9,8 @@ from sarracen import SarracenDataFrame -def _create_capture_pattern(def_int: Type[np.generic], - def_real: Type[np.generic], +def _create_capture_pattern(def_int: type[np.generic], + def_real: type[np.generic], swap_endian: bool = False) -> bytearray: """ Construct capture pattern. """ @@ -55,16 +54,16 @@ def _create_file_identifier(swap_endian: bool = False) -> bytearray: def _create_global_header(massoftype: float = 1e-6, - massoftype_7: Union[float, None] = None, - def_int: Type[np.generic] = np.int32, - def_real: Type[np.generic] = np.float64, + massoftype_7: float | None = None, + def_int: type[np.generic] = np.int32, + def_real: type[np.generic] = np.float64, mpi_blocks: int = 1, swap_endian: bool = False) -> bytearray: """ Construct global variables. Only massoftype in this example. """ dtypes = [def_int, np.int8, np.int16, np.int32, np.int64, def_real, np.float32, np.float64] - param_dicts: List[Dict] = [dict() for _ in dtypes] + param_dicts: list[dict] = [dict() for _ in dtypes] params_def_int = param_dicts[0] params_def_real = param_dicts[5] @@ -76,7 +75,7 @@ def _create_global_header(massoftype: float = 1e-6, params_def_int['nblocks'] = np.array([mpi_blocks], dtype=def_int) - dtype_param_pairs: List[Tuple[Type, Dict]] = list(zip(dtypes, param_dicts)) + dtype_param_pairs: list[tuple[type, dict]] = list(zip(dtypes, param_dicts)) read_tag = np.array([13], dtype='int32') if swap_endian: @@ -98,7 +97,7 @@ def _create_global_header(massoftype: float = 1e-6, file += bytearray(read_tag.tobytes()) for v in params.values(): - v_np = np.array([v], dtype=dtype) + v_np: np.ndarray = np.array([v], dtype=dtype) if swap_endian: v_np = v_np.byteswap() file += bytearray(v_np) @@ -109,7 +108,7 @@ def _create_global_header(massoftype: float = 1e-6, def _create_particle_array(tag: str, data: list, - dtype: Type[np.generic] = np.float64, + dtype: type[np.generic] = np.float64, swap_endian: bool = False) -> bytearray: read_tag = np.array([13], dtype='int32') @@ -136,8 +135,8 @@ def _create_particle_array(tag: str, (np.int32, np.float32, True), (np.int64, np.float64, True), (np.int64, np.float32, True)]) -def test_determine_default_precision(def_int: Type[np.generic], - def_real: Type[np.generic], +def test_determine_default_precision(def_int: type[np.generic], + def_real: type[np.generic], swap_endian: bool) -> None: """ Test if default int / real precision can be determined. """ diff --git a/sarracen/tests/test_render.py b/sarracen/tests/test_render.py index 2f7c864..6ee0197 100644 --- a/sarracen/tests/test_render.py +++ b/sarracen/tests/test_render.py @@ -1,5 +1,7 @@ """pytest unit tests for render.py functions.""" -from typing import Any, Callable, Dict, List +from collections.abc import Callable +from typing import Any + from matplotlib import pyplot as plt from numba import cuda from numpy.testing import assert_array_equal @@ -106,7 +108,7 @@ def test_cbar_exclusion(backend: str) -> None: sdf_3 = SarracenDataFrame(data_3) sdf_3.backend = backend - args_list: List[Dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, + args_list: list[dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, {'data': sdf_3, 'xsec': None}, {'data': sdf_3, 'xsec': 1.5}] for args in args_list: @@ -135,7 +137,7 @@ def test_cbar_keywords(backend: str) -> None: sdf_3 = SarracenDataFrame(data_3) sdf_3.backend = backend - args_list: List[Dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, + args_list: list[dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, {'data': sdf_3, 'xsec': None}, {'data': sdf_3, 'xsec': 1.5}] for args in args_list: @@ -167,7 +169,7 @@ def test_kwargs(backend: str) -> None: sdf_4 = SarracenDataFrame(data_4) sdf_4.backend = backend - args_list: List[Dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, + args_list: list[dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, {'data': sdf_3, 'xsec': None}, {'data': sdf_3, 'xsec': 1.5}] for args in args_list: @@ -216,7 +218,7 @@ def test_plot_labels(backend: str) -> None: sdf_3 = SarracenDataFrame(data_3) sdf_3.backend = backend - args_list: List[Dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, + args_list: list[dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, {'data': sdf_3, 'xsec': None}, {'data': sdf_3, 'xsec': 0}] for args in args_list: @@ -240,7 +242,7 @@ def test_plot_labels(backend: str) -> None: ('column ' if column else '') + 'rho' plt.close(fig) - functions_list: List[Callable] = [streamlines, arrowplot] + functions_list: list[Callable] = [streamlines, arrowplot] for func in functions_list: fig, ax = plt.subplots() func(sdf_2, ('Ax', 'Ay'), ax=ax) @@ -287,7 +289,7 @@ def test_plot_bounds(backend: str) -> None: sdf_3 = SarracenDataFrame(data_3) sdf_3.backend = backend - args_list: List[Dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, + args_list: list[dict[str, Any]] = [{'data': sdf_2, 'xsec': None}, {'data': sdf_3, 'xsec': None}, {'data': sdf_3, 'xsec': 1.5}] for args in args_list: @@ -319,7 +321,7 @@ def test_plot_bounds(backend: str) -> None: assert ax.get_ylim() == (0, interpolate_2d_line(sdf_2, 'P').max()) plt.close(fig) - functions_list: List[Callable] = [arrowplot, streamlines] + functions_list: list[Callable] = [arrowplot, streamlines] for func in functions_list: fig, ax = plt.subplots() func(sdf_3, ('Ax', 'Ay', 'Az'), ax=ax) diff --git a/sarracen/tests/writers/test_write_phantom.py b/sarracen/tests/writers/test_write_phantom.py index b1c3f0e..2e950e7 100644 --- a/sarracen/tests/writers/test_write_phantom.py +++ b/sarracen/tests/writers/test_write_phantom.py @@ -1,4 +1,4 @@ -from typing import Dict, Union, Any +from typing import Any import pandas as pd import numpy as np @@ -48,7 +48,7 @@ def test_check_for_essential_data(particles_df: pd.DataFrame) -> None: def test_params_dtype_standardization(particles_df: pd.DataFrame) -> None: """ Test that params dtype is correctly standardized.""" - params: Dict[str, Any] = {'massoftype': 1e-4, + params: dict[str, Any] = {'massoftype': 1e-4, 'ntypes': 8, 'nparttot': 8, 'massoftype_7': 1e-6, @@ -69,7 +69,7 @@ def test_validate_ntypes_12(particles_df: pd.DataFrame, id_method: str) -> None: """ Test that number of particle types is correctly validated.""" - params: Dict[str, Union[np.generic, str]] = {} + params: dict[str, np.generic | str] = {} if id_method == 'ntypes': params['ntypes'] = np.int32(12) @@ -100,7 +100,7 @@ def test_validate_ntypes_8(particles_df: pd.DataFrame, id_method: str) -> None: """ Test that number of particle types is correctly validated.""" - params: Dict[str, Union[np.generic, str]] = {} + params: dict[str, np.generic | str] = {} if id_method == 'ntypes': params['ntypes'] = np.int32(8) @@ -130,11 +130,9 @@ def test_validate_particle_counts(particles_df: pd.DataFrame, id_method: str) -> None: """ Test that number of each particle type is correctly autofilled.""" - params: Dict[str, - Union[np.generic, - str]] = {'massoftype': np.float64(1e-4), - 'ntypes': np.int32(8), - 'file_identifier': 'test Phantom write'} + params: dict[str, np.generic | str] = {'massoftype': np.float64(1e-4), + 'ntypes': np.int32(8), + 'file_identifier': 'test write'} # if npartoftype not present, then it determines based on itype alone if id_method == 'npartoftype': @@ -171,19 +169,18 @@ def test_validate_particle_counts(particles_df: pd.DataFrame, def test_params_reordering(particles_df: pd.DataFrame) -> None: """ Test that disordered params keyes are written correctly.""" - params: Dict[str, Union[np.generic, - str]] = {'massoftype': np.float64(1e-4), - 'massoftype_7': np.float64(1e-6), - 'massoftype_3': np.float64(1e-3), - 'nparttot': np.int64(8), - 'nparttot_2': np.int32(8), - 'npartoftype_15': np.int64(3), - 'npartoftype_3': np.int32(4), - 'npartoftype': np.int32(1), - 'npartoftype_7': np.int32(3), - 'npartoftype_11': np.int64(4), - 'npartoftype_9': np.int64(1), - 'file_identifier': 'test Phantom write'} + params: dict[str, np.generic | str] = {'massoftype': np.float64(1e-4), + 'massoftype_7': np.float64(1e-6), + 'massoftype_3': np.float64(1e-3), + 'nparttot': np.int64(8), + 'nparttot_2': np.int32(8), + 'npartoftype_15': np.int64(3), + 'npartoftype_3': np.int32(4), + 'npartoftype': np.int32(1), + 'npartoftype_7': np.int32(3), + 'npartoftype_11': np.int64(4), + 'npartoftype_9': np.int64(1), + 'file_identifier': 'test write'} params = _reorder_params(params) @@ -216,7 +213,7 @@ def test_validate_particle_mass_gas(particles_df: pd.DataFrame, id_method: bool) -> None: """ Test that mass of each particle type is correctly validated.""" - params: Dict[str, Union[np.generic, str]] = {'ntypes': np.int32(8)} + params: dict[str, np.generic | str] = {'ntypes': np.int32(8)} if id_method == 'params_mass': params['mass'] = np.float64(1e-4) @@ -253,7 +250,7 @@ def test_validate_particle_mass_dust(particles_df: pd.DataFrame, id_method: bool) -> None: """ Test that mass of each particle type is correctly validated.""" - params: Dict[str, Union[np.generic, str]] = {'ntypes': np.int32(8)} + params: dict[str, np.generic | str] = {'ntypes': np.int32(8)} if id_method == 'massoftype': params['massoftype'] = np.float64(1e-4) diff --git a/sarracen/writers/write_phantom.py b/sarracen/writers/write_phantom.py index 5d4bb8f..f0b0eb8 100644 --- a/sarracen/writers/write_phantom.py +++ b/sarracen/writers/write_phantom.py @@ -1,4 +1,5 @@ -from typing import Dict, List, Union, Type, Tuple, Any +from typing import Any, TypeAlias +from collections.abc import Sequence import numpy as np import re import warnings @@ -7,11 +8,11 @@ from ..sarracen_dataframe import SarracenDataFrame -params_dict_type = Dict[str, Union[np.generic, str]] +ParamsDict: TypeAlias = dict[str, np.generic | str] -def _write_fortran_block(value: List[np.generic], - dtype: Type[np.generic]) -> bytearray: +def _write_fortran_block(value: Sequence[int | float | np.generic], + dtype: type[np.generic]) -> bytearray: write_tag = np.array([len(value) * dtype().itemsize], dtype=np.int32) file = bytearray(write_tag.tobytes()) file += bytearray(np.array(value, dtype=dtype).tobytes()) @@ -46,8 +47,8 @@ def _write_file_identifier(sdf: SarracenDataFrame) -> bytearray: return file -def _write_capture_pattern(def_int: Type[np.generic], - def_real: Type[np.generic], +def _write_capture_pattern(def_int: type[np.generic], + def_real: type[np.generic], iversion: int = 1) -> bytearray: write_tag = np.array([16 + def_real().itemsize], dtype=np.int32) i1 = np.array([60769], dtype=def_int) @@ -72,7 +73,7 @@ def _check_for_essential_data(sdf: SarracenDataFrame) -> None: required_cols = [sdf.xcol, sdf.ycol, sdf.zcol, sdf.hcol] if any(col is None for col in required_cols): - missing: List[str] = [] + missing: list[str] = [] if sdf.xcol is None: missing.append("x") if sdf.ycol is None: @@ -112,7 +113,7 @@ def _check_for_essential_data(sdf: SarracenDataFrame) -> None: raise ValueError(msg) -def _remove_invalid_keys(params: params_dict_type) -> params_dict_type: +def _remove_invalid_keys(params: ParamsDict) -> ParamsDict: """ Remove keys specific to Sarracen internal use.""" exclude = ['file_identifier', 'mass', 'def_int_dtype', @@ -120,7 +121,7 @@ def _remove_invalid_keys(params: params_dict_type) -> params_dict_type: return {k: v for k, v in params.items() if k not in exclude} -def _standardize_dtypes(params: Dict[str, Any]) -> params_dict_type: +def _standardize_dtypes(params: dict[str, Any]) -> ParamsDict: """ Convert all params to numpy dtypes.""" for k, v in params.items(): @@ -133,7 +134,7 @@ def _standardize_dtypes(params: Dict[str, Any]) -> params_dict_type: def _validate_ntypes(sdf: SarracenDataFrame, - params: params_dict_type) -> params_dict_type: + params: ParamsDict) -> ParamsDict: """ Update params ntypes to reflect particle data.""" if 'ntypes' in params: @@ -183,7 +184,7 @@ def _validate_ntypes(sdf: SarracenDataFrame, def _validate_particle_counts(sdf: SarracenDataFrame, - params: params_dict_type) -> params_dict_type: + params: ParamsDict) -> ParamsDict: """ Update params particle counts to match actual particle counts. @@ -249,7 +250,7 @@ def _validate_particle_counts(sdf: SarracenDataFrame, def _validate_particle_masses(sdf: SarracenDataFrame, - params: params_dict_type) -> params_dict_type: + params: ParamsDict) -> ParamsDict: """ Update params particle masses to match actual particle masses.""" if 'mass' in params: @@ -287,7 +288,7 @@ def _validate_particle_masses(sdf: SarracenDataFrame, return params -def _relocate_special_params(param_dicts: List[params_dict_type]) -> None: +def _relocate_special_params(param_dicts: list[ParamsDict]) -> None: """Move certain parameters to their required dtype indices.""" # Integer parameters that must be int32 for param in ['iexternalforce', 'ieos']: @@ -309,12 +310,12 @@ def _rename_duplicate(tag: str) -> str: return tag -def _write_global_header_array(tags: List[str], - values: List[np.generic], - dtype: Type[np.generic]) -> bytearray: +def _write_global_header_array(tags: list[str], + values: list[np.generic], + dtype: type[np.generic]) -> bytearray: tags = [_rename_duplicate(tag) for tag in tags] - tags2: List[List[int]] = [list(map(ord, tag.ljust(16))) for tag in tags] - tags3: List[np.uint8] = [np.uint8(c) for tag in tags2 for c in tag] + tags2: list[list[int]] = [list(map(ord, tag.ljust(16))) for tag in tags] + tags3: list[np.uint8] = [np.uint8(c) for tag in tags2 for c in tag] file = _write_fortran_block(tags3, np.uint8) file += _write_fortran_block(values, dtype) @@ -322,7 +323,7 @@ def _write_global_header_array(tags: List[str], return file -def sort_key(k: str) -> Tuple[str, int]: +def sort_key(k: str) -> tuple[str, int]: """ Returns (basename, number) where basename is the part before any numeric suffix, and number is an integer used for sorting. @@ -337,13 +338,13 @@ def sort_key(k: str) -> Tuple[str, int]: return (base, num) -def _reorder_params(params: params_dict_type) -> params_dict_type: +def _reorder_params(params: ParamsDict) -> ParamsDict: return dict(sorted(params.items(), key=lambda item: sort_key(item[0]))) def _write_global_header(sdf: SarracenDataFrame, - def_int: Type[np.generic], - def_real: Type[np.generic]) -> bytearray: + def_int: type[np.generic], + def_real: type[np.generic]) -> bytearray: params_dict = sdf.params.copy() params_dict = _standardize_dtypes(params_dict) @@ -357,7 +358,7 @@ def _write_global_header(sdf: SarracenDataFrame, def_real, np.float32, np.float64] # create params dict per dtype and populate - param_dtype_dicts: List[params_dict_type] = [dict() for _ in dtypes] + param_dtype_dicts: list[ParamsDict] = [dict() for _ in dtypes] for k, v in params_dict.items(): if isinstance(v, def_int): param_dtype_dicts[0][k] = v @@ -406,7 +407,7 @@ def _write_global_header(sdf: SarracenDataFrame, def _get_array_tags(sdf: SarracenDataFrame, - dtype: Type[np.generic]) -> List[str]: + dtype: type[np.generic]) -> list[str]: return list(sdf.select_dtypes(include=[dtype]).columns) @@ -415,10 +416,9 @@ def _get_last_index(sdf: SarracenDataFrame) -> int: def _write_array_blocks(sdf: SarracenDataFrame, - def_int: Type[np.generic], - def_real: Type[np.generic], - sinks: Union[SarracenDataFrame, - None] = None) -> bytearray: + def_int: type[np.generic], + def_real: type[np.generic], + sinks: SarracenDataFrame | None = None) -> bytearray: dtypes = [def_int, np.int8, np.int16, np.int32, np.int64, def_real, np.float32, np.float64] @@ -460,9 +460,9 @@ def _write_array_blocks(sdf: SarracenDataFrame, # irregardless of what the xcol, ycol, etc are set to mapping = {sdf.xcol: 'x', sdf.ycol: 'y', sdf.zcol: 'z', sdf.hcol: 'h'} - write_tag = mapping.get(tag, tag) - write_tag = _rename_duplicate(write_tag).ljust(16) - file += _write_fortran_block(list(map(ord, write_tag)), + wtag = mapping.get(tag, tag) + wtag = _rename_duplicate(wtag).ljust(16) + file += _write_fortran_block(list(map(ord, wtag)), dtype=np.uint8) file += _write_fortran_block(list(sdf[tag]), dtype) return file @@ -470,7 +470,7 @@ def _write_array_blocks(sdf: SarracenDataFrame, def write_phantom(filename: str, sdf: SarracenDataFrame, - sinks: Union[SarracenDataFrame, None] = None) -> None: + sinks: SarracenDataFrame | None = None) -> None: """ Write a Phantom dump file from a SarracenDataFrame.