From 3abd8d9ffafd97395920218d3544c0c4c02ccefc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guillaume=20Daval-Fr=C3=A9rot?= Date: Fri, 10 Jan 2025 10:26:24 +0100 Subject: [PATCH] Clean documentation and add type hints (#228) * Clean documentation and add type hints * Apply ruff ANN rules to trajectories * Replace np.ndarray type hints with np.typing.NDArray * Fix missing import * Fix Literal use, change types to NDArray and reduce mypy errors * Add type alias for clustering coordinates --- src/mrinufft/trajectories/display.py | 130 ++++---- src/mrinufft/trajectories/gradients.py | 31 +- .../trajectories/inits/random_walk.py | 83 +++-- .../trajectories/inits/travelling_salesman.py | 99 +++--- src/mrinufft/trajectories/maths/fibonacci.py | 12 +- src/mrinufft/trajectories/maths/primes.py | 8 +- src/mrinufft/trajectories/maths/rotations.py | 39 ++- src/mrinufft/trajectories/maths/tsp_solver.py | 20 +- src/mrinufft/trajectories/sampling.py | 90 ++++-- src/mrinufft/trajectories/tools.py | 290 ++++++++++-------- src/mrinufft/trajectories/trajectory2D.py | 106 ++++--- src/mrinufft/trajectories/trajectory3D.py | 225 ++++++++------ src/mrinufft/trajectories/utils.py | 146 ++++----- 13 files changed, 732 insertions(+), 547 deletions(-) diff --git a/src/mrinufft/trajectories/display.py b/src/mrinufft/trajectories/display.py index 50f20344b..42e07cbe1 100644 --- a/src/mrinufft/trajectories/display.py +++ b/src/mrinufft/trajectories/display.py @@ -1,11 +1,15 @@ """Display functions for trajectories.""" +from __future__ import annotations + import itertools +from typing import Any import matplotlib as mpl import matplotlib.pyplot as plt import matplotlib.ticker as mticker import numpy as np +from numpy.typing import NDArray from .utils import ( DEFAULT_GMAX, @@ -46,7 +50,7 @@ class displayConfig: """Font size for most labels and texts, by default ``18``.""" small_fontsize: int = 14 """Font size for smaller texts, by default ``14``.""" - nb_colors = 10 + nb_colors: int = 10 """Number of colors to use in the color cycle, by default ``10``.""" palette: str = "tab10" """Name of the color palette to use, by default ``"tab10"``. @@ -58,33 +62,33 @@ class displayConfig: slewrate_point_color: str = "b" """Matplotlib color for slew rate constraint points, by default ``"b"`` (blue).""" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: # noqa ANN401 """Update the display configuration.""" self.update(**kwargs) - def update(self, **kwargs): + def update(self, **kwargs: Any) -> None: # noqa ANN401 """Update the display configuration.""" self._old_values = {} for key, value in kwargs.items(): self._old_values[key] = getattr(displayConfig, key) setattr(displayConfig, key, value) - def reset(self): + def reset(self) -> None: """Restore the display configuration.""" for key, value in self._old_values.items(): setattr(displayConfig, key, value) delattr(self, "_old_values") - def __enter__(self): + def __enter__(self) -> displayConfig: """Enter the context manager.""" return self - def __exit__(self, *args): + def __exit__(self, *args: Any) -> None: # noqa ANN401 """Exit the context manager.""" self.reset() @classmethod - def get_colorlist(cls): + def get_colorlist(cls) -> list[str | NDArray]: """Extract a list of colors from a matplotlib palette. If the palette is continuous, the colors will be sampled from it. @@ -124,7 +128,7 @@ def get_colorlist(cls): ############## -def _setup_2D_ticks(figsize, fig=None): +def _setup_2D_ticks(figsize: float, fig: plt.Figure | None = None) -> plt.Axes: """Add ticks to 2D plot.""" if fig is None: fig = plt.figure(figsize=(figsize, figsize)) @@ -139,7 +143,7 @@ def _setup_2D_ticks(figsize, fig=None): return ax -def _setup_3D_ticks(figsize, fig=None): +def _setup_3D_ticks(figsize: float, fig: plt.Figure | None = None) -> plt.Axes: """Add ticks to 3D plot.""" if fig is None: fig = plt.figure(figsize=(figsize, figsize)) @@ -163,21 +167,21 @@ def _setup_3D_ticks(figsize, fig=None): def display_2D_trajectory( - trajectory, - figsize=5, - one_shot=False, - subfigure=None, - show_constraints=False, - gmax=DEFAULT_GMAX, - smax=DEFAULT_SMAX, - constraints_order=None, - **constraints_kwargs, -): + trajectory: NDArray, + figsize: float = 5, + one_shot: bool | int = False, + subfigure: plt.Figure | plt.Axes | None = None, + show_constraints: bool = False, + gmax: float = DEFAULT_GMAX, + smax: float = DEFAULT_SMAX, + constraints_order: int | str | None = None, + **constraints_kwargs: Any, # noqa ANN401 +) -> plt.Axes: """Display 2D trajectories. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to display. figsize : float, optional Size of the figure. @@ -204,7 +208,7 @@ def display_2D_trajectory( typically 2 or `np.inf`, following the `numpy.linalg.norm` conventions on parameter `ord`. The default is None. - **kwargs + **constraints_kwargs Acquisition parameters used to check on hardware constraints, following the parameter convention from `mrinufft.trajectories.utils.compute_gradients_and_slew_rates`. @@ -278,23 +282,23 @@ def display_2D_trajectory( def display_3D_trajectory( - trajectory, - nb_repetitions=None, - figsize=5, - per_plane=True, - one_shot=False, - subfigure=None, - show_constraints=False, - gmax=DEFAULT_GMAX, - smax=DEFAULT_SMAX, - constraints_order=None, - **constraints_kwargs, -): + trajectory: NDArray, + nb_repetitions: int | None = None, + figsize: float = 5, + per_plane: bool = True, + one_shot: bool | int = False, + subfigure: plt.Figure | plt.Axes | None = None, + show_constraints: bool = False, + gmax: float = DEFAULT_GMAX, + smax: float = DEFAULT_SMAX, + constraints_order: int | str | None = None, + **constraints_kwargs: Any, # noqa ANN401 +) -> plt.Axes: """Display 3D trajectories. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to display. nb_repetitions : int Number of repetitions (planes, cones, shells, etc). @@ -417,22 +421,22 @@ def display_3D_trajectory( def display_gradients_simply( - trajectory, - shot_ids=(0,), - figsize=5, - fill_area=True, - show_signal=True, - uni_signal="gray", - uni_gradient=None, - subfigure=None, -): + trajectory: NDArray, + shot_ids: tuple[int, ...] = (0,), + figsize: float = 5, + fill_area: bool = True, + show_signal: bool = True, + uni_signal: str | None = "gray", + uni_gradient: str | None = None, + subfigure: plt.Figure | None = None, +) -> tuple[plt.Axes]: """Display gradients based on trajectory of any dimension. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to display. - shot_ids : list of int + shot_ids : tuple[int, ...], optional Indices of the shots to display. The default is `[0]`. figsize : float, optional @@ -455,7 +459,7 @@ def display_gradients_simply( unique color given as argument or just by the default color cycle when `None`. The default is `None`. - subfigure: plt.Figure or plt.SubFigure, optional + subfigure: plt.Figure, optional The figure where the trajectory should be displayed. The default is `None`. @@ -531,26 +535,26 @@ def display_gradients_simply( def display_gradients( - trajectory, - shot_ids=(0,), - figsize=5, - fill_area=True, - show_signal=True, - uni_signal="gray", - uni_gradient=None, - subfigure=None, - show_constraints=False, - gmax=DEFAULT_GMAX, - smax=DEFAULT_SMAX, - constraints_order=None, - raster_time=DEFAULT_RASTER_TIME, - **constraints_kwargs, -): + trajectory: NDArray, + shot_ids: tuple[int, ...] = (0,), + figsize: float = 5, + fill_area: bool = True, + show_signal: bool = True, + uni_signal: str | None = "gray", + uni_gradient: str | None = None, + subfigure: plt.Figure | plt.Axes | None = None, + show_constraints: bool = False, + gmax: float = DEFAULT_GMAX, + smax: float = DEFAULT_SMAX, + constraints_order: int | str | None = None, + raster_time: float = DEFAULT_RASTER_TIME, + **constraints_kwargs: Any, # noqa ANN401 +) -> tuple[plt.Axes]: """Display gradients based on trajectory of any dimension. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to display. shot_ids : list of int Indices of the shots to display. @@ -597,7 +601,7 @@ def display_gradients( Amount of time between the acquisition of two consecutive samples in ms. The default is `DEFAULT_RASTER_TIME`. - **kwargs + **constraints_kwargs Acquisition parameters used to check on hardware constraints, following the parameter convention from `mrinufft.trajectories.utils.compute_gradients_and_slew_rates`. diff --git a/src/mrinufft/trajectories/gradients.py b/src/mrinufft/trajectories/gradients.py index 39ac934c0..ac205cab1 100644 --- a/src/mrinufft/trajectories/gradients.py +++ b/src/mrinufft/trajectories/gradients.py @@ -1,24 +1,27 @@ """Functions to improve/modify gradients.""" +from typing import Callable + import numpy as np import numpy.linalg as nl +from numpy.typing import NDArray from scipy.interpolate import CubicSpline def patch_center_anomaly( - shot_or_params, - update_shot=None, - update_parameters=None, - in_out=False, - learning_rate=1e-1, -): + shot_or_params: NDArray | tuple, + update_shot: Callable[..., NDArray] | None = None, + update_parameters: Callable[..., tuple] | None = None, + in_out: bool = False, + learning_rate: float = 1e-1, +) -> tuple[NDArray, tuple]: """Re-position samples to avoid center anomalies. Some trajectories behave slightly differently from expected when approaching definition bounds, most often the k-space center as for spirals in some cases. - This function enforces non-strictly increasing monoticity of + This function enforces non-strictly increasing monotonicity of sample distances from the center, effectively reducing slew rates and smoothing gradient transitions locally. @@ -31,17 +34,17 @@ def patch_center_anomaly( shot_or_params : np.array, list Either a single shot of shape (Ns, Nd), or a list of arbitrary arguments used by ``update_shot`` to initialize a single shot. - update_shot : function, optional + update_shot : Callable[..., NDArray], optional Function used to initialize a single shot based on parameters provided by ``update_parameters``. If None, cubic splines are used as an approximation instead, by default None - update_parameters : function, optional + update_parameters : Callable[..., tuple], optional Function used to update shot parameters when provided in ``shot_or_params`` from an updated shot and parameters. If None, cubic spline parameterization is used instead, by default None in_out : bool, optional - Whether the shot is going in-and-out or start from the center, + Whether the shot is going in-and-out or starts from the center, by default False learning_rate : float, optional Learning rate used in the iterative optimization process, @@ -49,9 +52,9 @@ def patch_center_anomaly( Returns ------- - array_like + NDArray N-D trajectory based on ``shot_or_params`` if a shot or - update_shot otherwise. + ``update_shot`` otherwise. list Updated parameters either in the ``shot_or_params`` format if params, or cubic spline parameterization as an array of @@ -70,7 +73,7 @@ def patch_center_anomaly( if update_shot is None or update_parameters is None: - def _default_update_parameters(shot, *parameters): + def _default_update_parameters(shot: NDArray, *parameters: list) -> list: return parameters update_parameters = _default_update_parameters @@ -114,5 +117,5 @@ def _default_update_parameters(shot, *parameters): single_shot = cbs(x_axis).T parameters = update_parameters(single_shot, *parameters) - single_shot = single_shot = update_shot(*parameters) + single_shot = update_shot(*parameters) return single_shot, parameters diff --git a/src/mrinufft/trajectories/inits/random_walk.py b/src/mrinufft/trajectories/inits/random_walk.py index da68feb1f..889fd20dd 100644 --- a/src/mrinufft/trajectories/inits/random_walk.py +++ b/src/mrinufft/trajectories/inits/random_walk.py @@ -1,17 +1,19 @@ """Trajectories based on random walks.""" +from typing import Any, Literal + import numpy as np +from numpy.typing import NDArray from ..sampling import sample_from_density from ..utils import KMAX -def _get_adjacent_neighbors_offsets(shape): - return np.concatenate([np.eye(len(shape)), -np.eye(len(shape))], axis=0).astype(int) +def _get_adjacent_neighbors_offsets(nb_dims: int) -> NDArray: + return np.concatenate([np.eye(nb_dims), -np.eye(nb_dims)], axis=0).astype(int) -def _get_neighbors_offsets(shape): - nb_dims = len(shape) +def _get_neighbors_offsets(nb_dims: int) -> NDArray: neighbors = (np.indices([3] * nb_dims) - 1).reshape((nb_dims, -1)).T nb_half = neighbors.shape[0] // 2 # Remove full zero entry @@ -20,18 +22,25 @@ def _get_neighbors_offsets(shape): def _initialize_ND_random_walk( - Nc, Ns, density, *, diagonals=True, pseudo_random=True, **sampling_kwargs -): + Nc: int, + Ns: int, + density: NDArray, + *, + diagonals: bool = True, + pseudo_random: bool = True, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: density = density / np.sum(density) flat_density = np.copy(density.flatten()) - shape = np.array(density.shape) + shape = density.shape + nb_dims = len(shape) mask = np.ones_like(flat_density) # Prepare neighbor offsets once offsets = ( - _get_neighbors_offsets(shape) + _get_neighbors_offsets(nb_dims) if diagonals - else _get_adjacent_neighbors_offsets(shape) + else _get_adjacent_neighbors_offsets(nb_dims) ) # Make all random draws at once for performance @@ -39,9 +48,8 @@ def _initialize_ND_random_walk( # Initialize shot starting points locations = sample_from_density(Nc, density, **sampling_kwargs) - choices = np.around((locations + KMAX) * (np.array(density.shape) - 1)).astype(int) - choices = np.ravel_multi_index(choices.T, density.shape) - # choices = np.random.choice(np.arange(len(flat_density)), size=Nc, p=flat_density) + choices = np.around((locations + KMAX) * (np.array(shape) - 1)).astype(int) + choices = np.ravel_multi_index(choices.T, shape) routes = [choices] # Walk @@ -52,7 +60,7 @@ def _initialize_ND_random_walk( # Find out-of-bound neighbors and ignore them invalids = (neighbors < 0).any(axis=0) | ( - neighbors >= shape[:, None, None] + neighbors >= np.array(shape)[:, None, None] ).any(axis=0) neighbors[:, invalids] = 0 invalids = invalids.T @@ -71,25 +79,30 @@ def _initialize_ND_random_walk( choices = neighbors[np.arange(Nc), indices] routes.append(choices) - # Update density to account for already drawed positions + # Update density to account for already drawn positions if pseudo_random: flat_density[choices] = ( mask[choices] * flat_density[choices] / (mask[choices] + 1) ) mask[choices] += 1 - routes = np.array(routes).T # Create trajectory from routes locations = np.indices(shape) - locations = locations.reshape((len(shape), -1)) - trajectory = np.array([locations[:, r].T for r in routes]) - trajectory = 2 * KMAX * trajectory / (shape - 1) - KMAX + locations = locations.reshape((nb_dims, -1)) + trajectory = np.array([locations[:, r].T for r in np.array(routes).T]) + trajectory = 2 * KMAX * trajectory / (np.array(shape) - 1) - KMAX return trajectory def initialize_2D_random_walk( - Nc, Ns, density, *, diagonals=True, pseudo_random=True, **sampling_kwargs -): + Nc: int, + Ns: int, + density: NDArray, + *, + diagonals: bool = True, + pseudo_random: bool = True, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: """Initialize a 2D random walk trajectory. This is an adaptation of the proposition from [Cha+14]_. @@ -107,31 +120,31 @@ def initialize_2D_random_walk( Number of shots Ns : int Number of samples per shot - density : array_like + density : NDArray Sampling density used to determine the walk probabilities, normalized automatically by its sum during the call for convenience. diagonals : bool, optional - Whether to draw the next walk step from the diagional neighbors + Whether to draw the next walk step from the diagonal neighbors on top of the adjacent ones. Default to ``True``. pseudo_random : bool, optional Whether to adapt the density dynamically to reduce areas already covered. The density is still statistically followed for undersampled acquisitions. Default to ``True``. - **sampling_kwargs + **sampling_kwargs : Any Sampling parameters in ``mrinufft.trajectories.sampling.sample_from_density`` used for the shot starting positions. Returns ------- - array_like + NDArray 2D random walk trajectory References ---------- .. [Cha+14] Chauffert, Nicolas, Philippe Ciuciu, Jonas Kahn, and Pierre Weiss. - "Variable density sampling with continuous trajectories." + "Variable density sampling with continuous trajectories" SIAM Journal on Imaging Sciences 7, no. 4 (2014): 1962-1992. """ if len(density.shape) != 2: @@ -147,8 +160,14 @@ def initialize_2D_random_walk( def initialize_3D_random_walk( - Nc, Ns, density, *, diagonals=True, pseudo_random=True, **sampling_kwargs -): + Nc: int, + Ns: int, + density: NDArray, + *, + diagonals: bool = True, + pseudo_random: bool = True, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: """Initialize a 3D random walk trajectory. This is an adaptation of the proposition from [Cha+14]_. @@ -166,31 +185,31 @@ def initialize_3D_random_walk( Number of shots Ns : int Number of samples per shot - density : array_like + density : NDArray Sampling density used to determine the walk probabilities, normalized automatically by its sum during the call for convenience. diagonals : bool, optional - Whether to draw the next walk step from the diagional neighbors + Whether to draw the next walk step from the diagonal neighbors on top of the adjacent ones. Default to ``True``. pseudo_random : bool, optional Whether to adapt the density dynamically to reduce areas already covered. The density is still statistically followed for undersampled acquisitions. Default to ``True``. - **sampling_kwargs + **sampling_kwargs : Any Sampling parameters in ``mrinufft.trajectories.sampling.sample_from_density`` used for the shot starting positions. Returns ------- - array_like + NDArray 3D random walk trajectory References ---------- .. [Cha+14] Chauffert, Nicolas, Philippe Ciuciu, Jonas Kahn, and Pierre Weiss. - "Variable density sampling with continuous trajectories." + "Variable density sampling with continuous trajectories" SIAM Journal on Imaging Sciences 7, no. 4 (2014): 1962-1992. """ if len(density.shape) != 3: diff --git a/src/mrinufft/trajectories/inits/travelling_salesman.py b/src/mrinufft/trajectories/inits/travelling_salesman.py index 6b4f7764b..9e3c35762 100644 --- a/src/mrinufft/trajectories/inits/travelling_salesman.py +++ b/src/mrinufft/trajectories/inits/travelling_salesman.py @@ -1,7 +1,10 @@ """Trajectories based on the Travelling Salesman Problem.""" +from typing import Any, Literal, TypeAlias + import numpy as np import numpy.linalg as nl +from numpy.typing import NDArray from scipy.interpolate import CubicSpline from tqdm.auto import tqdm @@ -9,8 +12,10 @@ from ..sampling import sample_from_density from ..tools import oversample +Coordinate: TypeAlias = Literal["x", "y", "z", "r", "phi", "theta"] + -def _get_approx_cluster_sizes(nb_total, nb_clusters): +def _get_approx_cluster_sizes(nb_total: int, nb_clusters: int) -> NDArray: # Give a list of cluster sizes close to sqrt(`nb_total`) cluster_sizes = round(nb_total / nb_clusters) * np.ones(nb_clusters).astype(int) delta_sum = nb_total - np.sum(cluster_sizes) @@ -18,7 +23,7 @@ def _get_approx_cluster_sizes(nb_total, nb_clusters): return cluster_sizes -def _sort_by_coordinate(array, coord): +def _sort_by_coordinate(array: NDArray, coord: Coordinate) -> NDArray: # Sort a list of N-D locations by a Cartesian/spherical coordinate if array.shape[-1] < 3 and coord.lower() in ["z", "theta"]: raise ValueError( @@ -47,8 +52,12 @@ def _sort_by_coordinate(array, coord): def _cluster_by_coordinate( - locations, nb_clusters, cluster_by, second_cluster_by=None, sort_by=None -): + locations: NDArray, + nb_clusters: int, + cluster_by: Coordinate, + second_cluster_by: Coordinate | None = None, + sort_by: Coordinate | None = None, +) -> NDArray: # Cluster approximately a list of N-D locations by Cartesian/spherical coordinates # Gather dimension variables nb_dims = locations.shape[-1] @@ -87,17 +96,17 @@ def _cluster_by_coordinate( def _initialize_ND_travelling_salesman( - Nc, - Ns, - density, - first_cluster_by=None, - second_cluster_by=None, - sort_by=None, - tsp_tol=1e-8, + Nc: int, + Ns: int, + density: NDArray, + first_cluster_by: Coordinate | None = None, + second_cluster_by: Coordinate | None = None, + sort_by: Coordinate | None = None, + tsp_tol: float = 1e-8, *, - verbose=False, - **sampling_kwargs, -): + verbose: bool = False, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: # Check arguments validity if Nc * Ns > np.prod(density.shape): raise ValueError("`density` array not large enough to pick `Nc` * `Ns` points.") @@ -134,17 +143,17 @@ def _initialize_ND_travelling_salesman( def initialize_2D_travelling_salesman( - Nc, - Ns, - density, - first_cluster_by=None, - second_cluster_by=None, - sort_by=None, - tsp_tol=1e-8, + Nc: int, + Ns: int, + density: NDArray, + first_cluster_by: Coordinate | None = None, + second_cluster_by: Coordinate | None = None, + sort_by: Coordinate | None = None, + tsp_tol: float = 1e-8, *, - verbose=False, - **sampling_kwargs, -): + verbose: bool = False, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: """ Initialize a 2D trajectory using a Travelling Salesman Problem (TSP)-based path. @@ -160,14 +169,14 @@ def initialize_2D_travelling_salesman( The number of clusters (or shots) to divide the trajectory into. Ns : int The number of points per cluster. - density : np.ndarray + density : NDArray A 2-dimensional density array from which points are sampled. - first_cluster_by : str, optional + first_cluster_by : {"x", "y", "z", "r", "phi", "theta"}, optional The coordinate used to cluster points initially, by default ``None``. - second_cluster_by : str, optional + second_cluster_by : {"x", "y", "z", "r", "phi", "theta"}, optional A secondary coordinate used for clustering within primary clusters, by default ``None``. - sort_by : str, optional + sort_by : {"x", "y", "z", "r", "phi", "theta"}, optional The coordinate by which to order points within each cluster, by default ``None``. tsp_tol : float, optional @@ -180,7 +189,7 @@ def initialize_2D_travelling_salesman( Returns ------- - np.ndarray + NDArray A 2D array representing the TSP-ordered trajectory. Raises @@ -192,7 +201,7 @@ def initialize_2D_travelling_salesman( ---------- .. [Cha+14] Chauffert, Nicolas, Philippe Ciuciu, Jonas Kahn, and Pierre Weiss. - "Variable density sampling with continuous trajectories." + "Variable density sampling with continuous trajectories" SIAM Journal on Imaging Sciences 7, no. 4 (2014): 1962-1992. """ if len(density.shape) != 2: @@ -211,17 +220,17 @@ def initialize_2D_travelling_salesman( def initialize_3D_travelling_salesman( - Nc, - Ns, - density, - first_cluster_by=None, - second_cluster_by=None, - sort_by=None, - tsp_tol=1e-8, + Nc: int, + Ns: int, + density: NDArray, + first_cluster_by: Coordinate | None = None, + second_cluster_by: Coordinate | None = None, + sort_by: Coordinate | None = None, + tsp_tol: float = 1e-8, *, - verbose=False, - **sampling_kwargs, -): + verbose: bool = False, + **sampling_kwargs: Any, # noqa ANN401 +) -> NDArray: """ Initialize a 3D trajectory using a Travelling Salesman Problem (TSP)-based path. @@ -239,14 +248,14 @@ def initialize_3D_travelling_salesman( The number of clusters (or shots) to divide the trajectory into. Ns : int The number of points per cluster. - density : np.ndarray + density : NDArray A 3-dimensional density array from which points are sampled. - first_cluster_by : str, optional + first_cluster_by : {"x", "y", "z", "r", "phi", "theta"}, optional The coordinate used to cluster points initially, by default ``None``. - second_cluster_by : str, optional + second_cluster_by : {"x", "y", "z", "r", "phi", "theta"}, optional A secondary coordinate used for clustering within primary clusters, by default ``None``. - sort_by : str, optional + sort_by : {"x", "y", "z", "r", "phi", "theta"}, optional The coordinate by which to order points within each cluster, by default ``None``. tsp_tol : float, optional @@ -259,7 +268,7 @@ def initialize_3D_travelling_salesman( Returns ------- - np.ndarray + NDArray A 3D array representing the TSP-ordered trajectory. Raises diff --git a/src/mrinufft/trajectories/maths/fibonacci.py b/src/mrinufft/trajectories/maths/fibonacci.py index d7cc20a88..69e621210 100644 --- a/src/mrinufft/trajectories/maths/fibonacci.py +++ b/src/mrinufft/trajectories/maths/fibonacci.py @@ -3,7 +3,7 @@ import numpy as np -def is_from_fibonacci_sequence(n): +def is_from_fibonacci_sequence(n: int) -> bool: """Check if an integer belongs to the Fibonacci sequence. An integer belongs to the Fibonacci sequence if either @@ -21,14 +21,14 @@ def is_from_fibonacci_sequence(n): Whether or not ``n`` belongs to the Fibonacci sequence. """ - def _is_perfect_square(n): + def _is_perfect_square(n: int) -> bool: r = int(np.sqrt(n)) return r * r == n return _is_perfect_square(5 * n**2 + 4) or _is_perfect_square(5 * n**2 - 4) -def get_closest_fibonacci_number(x): +def get_closest_fibonacci_number(x: float) -> int: """Provide the closest Fibonacci number. Parameters @@ -52,7 +52,7 @@ def get_closest_fibonacci_number(x): return xf -def generate_fibonacci_lattice(nb_points, epsilon=0.25): +def generate_fibonacci_lattice(nb_points: int, epsilon: float = 0.25) -> np.ndarray: """Generate 2D Cartesian coordinates using the Fibonacci lattice. Place 2D points over a 1x1 square following the Fibonacci lattice. @@ -78,7 +78,7 @@ def generate_fibonacci_lattice(nb_points, epsilon=0.25): return fibonacci_square -def generate_fibonacci_circle(nb_points, epsilon=0.25): +def generate_fibonacci_circle(nb_points: int, epsilon: float = 0.25) -> np.ndarray: """Generate 2D Cartesian coordinates shaped as Fibonacci spirals. Place 2D points structured as Fibonacci spirals by distorting @@ -106,7 +106,7 @@ def generate_fibonacci_circle(nb_points, epsilon=0.25): return fibonacci_circle -def generate_fibonacci_sphere(nb_points, epsilon=0.25): +def generate_fibonacci_sphere(nb_points: int, epsilon: float = 0.25) -> np.ndarray: """Generate 3D Cartesian coordinates as a Fibonacci sphere. Place 3D points almost evenly over a sphere surface of radius diff --git a/src/mrinufft/trajectories/maths/primes.py b/src/mrinufft/trajectories/maths/primes.py index 3a9df50cb..9db01c483 100644 --- a/src/mrinufft/trajectories/maths/primes.py +++ b/src/mrinufft/trajectories/maths/primes.py @@ -3,7 +3,9 @@ import numpy as np -def compute_coprime_factors(Nc, length, start=1, update=1): +def compute_coprime_factors( + Nc: int, length: int, start: int = 1, update: int = 1 +) -> list[int]: """Compute a list of coprime factors of Nc. Parameters @@ -19,11 +21,11 @@ def compute_coprime_factors(Nc, length, start=1, update=1): Returns ------- - list + list[int] List of coprime factors of Nc. """ count = start - coprimes = [] + coprimes: list[int] = [] while len(coprimes) < length: # Check greatest common divider (gcd) if np.gcd(Nc, count) == 1: diff --git a/src/mrinufft/trajectories/maths/rotations.py b/src/mrinufft/trajectories/maths/rotations.py index 568780786..f91cd37b8 100644 --- a/src/mrinufft/trajectories/maths/rotations.py +++ b/src/mrinufft/trajectories/maths/rotations.py @@ -2,9 +2,10 @@ import numpy as np import numpy.linalg as nl +from numpy.typing import NDArray -def R2D(theta): +def R2D(theta: float) -> NDArray: """Initialize 2D rotation matrix. Parameters @@ -14,13 +15,13 @@ def R2D(theta): Returns ------- - np.ndarray + NDArray 2D rotation matrix. """ return np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]) -def Rx(theta): +def Rx(theta: float) -> NDArray: """Initialize 3D rotation matrix around x axis. Parameters @@ -30,7 +31,7 @@ def Rx(theta): Returns ------- - np.ndarray + NDArray 3D rotation matrix. """ return np.array( @@ -42,7 +43,7 @@ def Rx(theta): ) -def Ry(theta): +def Ry(theta: float) -> NDArray: """Initialize 3D rotation matrix around y axis. Parameters @@ -52,7 +53,7 @@ def Ry(theta): Returns ------- - np.ndarray + NDArray 3D rotation matrix. """ return np.array( @@ -64,7 +65,7 @@ def Ry(theta): ) -def Rz(theta): +def Rz(theta: float) -> NDArray: """Initialize 3D rotation matrix around z axis. Parameters @@ -74,7 +75,7 @@ def Rz(theta): Returns ------- - np.ndarray + NDArray 3D rotation matrix. """ return np.array( @@ -86,7 +87,13 @@ def Rz(theta): ) -def Rv(v1, v2, normalize=True, eps=1e-8): +def Rv( + v1: NDArray, + v2: NDArray, + eps: float = 1e-8, + *, + normalize: bool = True, +) -> NDArray: """Initialize 3D rotation matrix from two vectors. Initialize a 3D rotation matrix from two vectors using Rodrigues's rotation @@ -97,16 +104,18 @@ def Rv(v1, v2, normalize=True, eps=1e-8): Parameters ---------- - v1 : np.ndarray + v1 : NDArray Source vector. - v2 : np.ndarray + v2 : NDArray Target vector. + eps : float, optional + Tolerance to consider two vectors as colinear. The default is 1e-8. normalize : bool, optional Normalize the vectors. The default is True. Returns ------- - np.ndarray + NDArray 3D rotation matrix. """ # Check for colinearity, not handled by Rodrigues' coefficients @@ -122,7 +131,7 @@ def Rv(v1, v2, normalize=True, eps=1e-8): return np.identity(3) + cross_matrix + cross_matrix @ cross_matrix / (1 + cos_theta) -def Ra(vector, theta): +def Ra(vector: NDArray, theta: float) -> NDArray: """Initialize 3D rotation matrix around an arbitrary vector. Initialize a 3D rotation matrix to rotate around `vector` by an angle `theta`. @@ -130,14 +139,14 @@ def Ra(vector, theta): Parameters ---------- - vector : np.ndarray + vector : NDArray Vector defining the rotation axis, automatically normalized. theta : float Angle in radians defining the rotation around `vector`. Returns ------- - np.ndarray + NDArray 3D rotation matrix. """ cos_t = np.cos(theta) diff --git a/src/mrinufft/trajectories/maths/tsp_solver.py b/src/mrinufft/trajectories/maths/tsp_solver.py index b4bcf859d..0b40a53aa 100644 --- a/src/mrinufft/trajectories/maths/tsp_solver.py +++ b/src/mrinufft/trajectories/maths/tsp_solver.py @@ -1,29 +1,31 @@ """Solver for the Travelling Salesman Problem.""" import numpy as np +from numpy.typing import NDArray -def solve_tsp_with_2opt(locations, improvement_threshold=1e-8): +def solve_tsp_with_2opt( + locations: NDArray, improvement_threshold: float = 1e-8 +) -> NDArray: """Solve the TSP problem using a 2-opt approach. A sub-optimal solution to the Travelling Salesman Problem (TSP) is provided using the 2-opt approach in O(n²) where chunks of an initially random route are reversed, and selected if the - total distance is reduced. As a result the route solution + total distance is reduced. As a result, the route solution does not cross its own path in 2D. Parameters ---------- - locations : array_like - An array of N points with shape (N, D) with D - the space dimension. - improvement_threshold: float - Threshold used as progress criterion to stop the optimization - process. + locations : NDArray + An array of N points with shape (N, D) where D is the space dimension. + improvement_threshold : float, optional + Threshold used as progress criterion to stop the optimization process. + The default is 1e-8. Returns ------- - array_like + NDArray The new positions order of shape (N,). """ route = np.arange(locations.shape[0]) diff --git a/src/mrinufft/trajectories/sampling.py b/src/mrinufft/trajectories/sampling.py index 14ed18fc0..6b4c643ad 100644 --- a/src/mrinufft/trajectories/sampling.py +++ b/src/mrinufft/trajectories/sampling.py @@ -1,17 +1,27 @@ """Sampling densities and methods.""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + import pywt as pw + import numpy as np import numpy.fft as nf -import numpy.linalg as nl -import numpy.random as nr +from numpy.typing import NDArray from tqdm.auto import tqdm from .utils import KMAX def sample_from_density( - nb_samples, density, method="random", *, dim_compensation="auto" -): + nb_samples: int, + density: NDArray, + method: Literal["random", "lloyd"] = "random", + *, + dim_compensation: Literal["auto"] | bool = "auto", +) -> NDArray: """ Sample points based on a given density distribution. @@ -19,14 +29,14 @@ def sample_from_density( ---------- nb_samples : int The number of samples to draw. - density : np.ndarray + density : NDArray An array representing the density distribution from which samples are drawn, normalized automatically by its sum during the call for convenience. - method : str, optional + method : Literal["random", "lloyd"], optional The sampling method to use, either 'random' for random sampling over the discrete grid defined by the density or 'lloyd' for Lloyd's algorithm over a continuous space, by default "random". - dim_compensation : str, bool, optional + dim_compensation : Literal["auto"], bool, optional Whether to apply a specific dimensionality compensation introduced in [Cha+14]_. An exponent ``N/(N-1)`` with ``N`` the number of dimensions in ``density`` is applied to fix the observed @@ -37,7 +47,7 @@ def sample_from_density( Returns ------- - np.ndarray + NDArray An array of range-normalized sampled locations. Raises @@ -50,7 +60,7 @@ def sample_from_density( ---------- .. [Cha+14] Chauffert, Nicolas, Philippe Ciuciu, Jonas Kahn, and Pierre Weiss. - "Variable density sampling with continuous trajectories." + "Variable density sampling with continuous trajectories" SIAM Journal on Imaging Sciences 7, no. 4 (2014): 1962-1992. """ try: @@ -63,7 +73,7 @@ def sample_from_density( ) from err # Define dimension variables - shape = np.array(density.shape) + shape = density.shape nb_dims = len(shape) max_nb_samples = np.prod(shape) density = density / np.sum(density) @@ -81,7 +91,7 @@ def sample_from_density( density = density / np.sum(density) # Sample using specified method - rng = nr.default_rng() + rng = np.random.default_rng() if method == "random": choices = rng.choice( np.arange(max_nb_samples), @@ -91,7 +101,7 @@ def sample_from_density( ) locations = np.indices(shape).reshape((nb_dims, -1))[:, choices] locations = locations.T + 0.5 - locations = locations / shape[None, :] + locations = locations / np.array(shape)[None, :] locations = 2 * KMAX * locations - KMAX elif method == "lloyd": kmeans = ( @@ -100,17 +110,22 @@ def sample_from_density( else BisectingKMeans(n_clusters=nb_samples) ) kmeans.fit( - np.indices(density.shape).reshape((nb_dims, -1)).T, + np.indices(shape).reshape((nb_dims, -1)).T, sample_weight=density.flatten(), ) - locations = kmeans.cluster_centers_ - np.array(density.shape) / 2 + locations = kmeans.cluster_centers_ - np.array(shape) / 2 locations = KMAX * locations / np.max(np.abs(locations)) else: raise ValueError(f"Unknown sampling method {method}.") return locations -def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): +def create_cutoff_decay_density( + shape: tuple[int, ...], + cutoff: float, + decay: float, + resolution: NDArray | None = None, +) -> NDArray: """ Create a density with central plateau and polynomial decay. @@ -120,7 +135,7 @@ def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): Parameters ---------- - shape : tuple of int + shape : tuple[int, ...] The shape of the density grid, analog to the field-of-view as opposed to ``resolution`` below. cutoff : float @@ -128,13 +143,13 @@ def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): and 1 within which density remains uniform and beyond which it decays. decay : float The polynomial decay in density beyond the cutoff ratio. - resolution : np.ndarray, optional + resolution : NDArray, optional Resolution scaling factors for each dimension of the density grid, by default ``None``. Returns ------- - np.ndarray + NDArray A density array with values decaying based on the specified cutoff ratio and decay rate. @@ -146,7 +161,6 @@ def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): magnetic resonance imaging." IEEE Transactions on Medical Imaging 41, no. 8 (2022): 2105-2117. """ - shape = np.array(shape) nb_dims = len(shape) if not resolution: @@ -156,7 +170,7 @@ def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): for i in range(nb_dims): differences[i] = differences[i] + 0.5 - shape[i] / 2 differences[i] = differences[i] / shape[i] / resolution[i] - distances = nl.norm(differences, axis=0) + distances = np.linalg.norm(differences, axis=0) cutoff = cutoff * np.max(differences) if cutoff else np.min(differences) density = np.ones(shape) @@ -167,7 +181,9 @@ def create_cutoff_decay_density(shape, cutoff, decay, resolution=None): return density -def create_polynomial_density(shape, decay, resolution=None): +def create_polynomial_density( + shape: tuple[int, ...], decay: float, resolution: NDArray | None = None +) -> NDArray: """ Create a density with polynomial decay from the center. @@ -177,13 +193,13 @@ def create_polynomial_density(shape, decay, resolution=None): The shape of the density grid. decay : float The exponent that controls the rate of decay for density. - resolution : np.ndarray, optional + resolution : NDArray, optional Resolution scaling factors for each dimension of the density grid, by default None. Returns ------- - np.ndarray + NDArray A density array with polynomial decay. """ return create_cutoff_decay_density( @@ -191,7 +207,7 @@ def create_polynomial_density(shape, decay, resolution=None): ) -def create_energy_density(dataset): +def create_energy_density(dataset: NDArray) -> NDArray: """ Create a density based on energy in the Fourier spectrum. @@ -200,7 +216,7 @@ def create_energy_density(dataset): Parameters ---------- - dataset : np.ndarray + dataset : NDArray The dataset from which to calculate the density based on its Fourier transform, with an expected shape (nb_volumes, dim_1, ..., dim_N). @@ -208,7 +224,7 @@ def create_energy_density(dataset): Returns ------- - np.ndarray + NDArray A density array derived from the mean energy in the Fourier domain of the input dataset. """ @@ -221,7 +237,13 @@ def create_energy_density(dataset): return density -def create_chauffert_density(shape, wavelet_basis, nb_wavelet_scales, verbose=False): +def create_chauffert_density( + shape: tuple[int, ...], + wavelet_basis: str | pw.Wavelet, + nb_wavelet_scales: int, + *, + verbose: bool = False, +) -> NDArray: """Create a density based on Chauffert's method. This is a reproduction of the proposition from [CCW13]_. @@ -231,7 +253,7 @@ def create_chauffert_density(shape, wavelet_basis, nb_wavelet_scales, verbose=Fa Parameters ---------- - shape : tuple of int + shape : tuple[int, ...] The shape of the density grid. wavelet_basis : str, pywt.Wavelet The wavelet basis to use for wavelet decomposition, either @@ -244,7 +266,7 @@ def create_chauffert_density(shape, wavelet_basis, nb_wavelet_scales, verbose=Fa Returns ------- - np.ndarray + NDArray A density array created based on wavelet transform coefficients. See Also @@ -290,7 +312,11 @@ def create_chauffert_density(shape, wavelet_basis, nb_wavelet_scales, verbose=Fa return nf.ifftshift(density) -def create_fast_chauffert_density(shape, wavelet_basis, nb_wavelet_scales): +def create_fast_chauffert_density( + shape: tuple[int, ...], + wavelet_basis: str | pw.Wavelet, + nb_wavelet_scales: int, +) -> NDArray: """Create a density based on an approximated Chauffert method. This implementation is based on this @@ -306,7 +332,7 @@ def create_fast_chauffert_density(shape, wavelet_basis, nb_wavelet_scales): Parameters ---------- - shape : tuple of int + shape : tuple[int, ...] The shape of the density grid. wavelet_basis : str, pywt.Wavelet The wavelet basis to use for wavelet decomposition, either @@ -317,7 +343,7 @@ def create_fast_chauffert_density(shape, wavelet_basis, nb_wavelet_scales): Returns ------- - np.ndarray + NDArray A density array created using a faster approximation based on 1D projections of the wavelet transform. diff --git a/src/mrinufft/trajectories/tools.py b/src/mrinufft/trajectories/tools.py index cc099c07a..02ca23ae0 100644 --- a/src/mrinufft/trajectories/tools.py +++ b/src/mrinufft/trajectories/tools.py @@ -1,6 +1,9 @@ """Functions to manipulate/modify trajectories.""" +from typing import Any, Callable, Literal + import numpy as np +from numpy.typing import NDArray from scipy.interpolate import CubicSpline, interp1d from .maths import Rv, Rx, Ry, Rz @@ -11,23 +14,29 @@ ################ -def stack(trajectory, nb_stacks, z_tilt=None, hard_bounded=True): +def stack( + trajectory: NDArray, + nb_stacks: int, + z_tilt: str | float | None = None, + *, + hard_bounded: bool = True, +) -> NDArray: """Stack 2D or 3D trajectories over the :math:`k_z`-axis. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory in 2D or 3D to stack. nb_stacks : int Number of stacks repeating the provided trajectory. - z_tilt : str, optional + z_tilt : str | float, optional Tilt of the stacks, by default `None`. hard_bounded : bool, optional Whether the stacks should be strictly within the limits of the k-space. Returns ------- - array_like + NDArray Stacked trajectory. """ # Check dimensionality and initialize output @@ -55,25 +64,31 @@ def stack(trajectory, nb_stacks, z_tilt=None, hard_bounded=True): return new_trajectory.reshape(nb_stacks * Nc, Ns, 3) -def rotate(trajectory, nb_rotations, x_tilt=None, y_tilt=None, z_tilt=None): +def rotate( + trajectory: NDArray, + nb_rotations: int, + x_tilt: str | float | None = None, + y_tilt: str | float | None = None, + z_tilt: str | float | None = None, +) -> NDArray: """Rotate 2D or 3D trajectories over the different axes. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory in 2D or 3D to rotate. nb_rotations : int Number of rotations repeating the provided trajectory. - x_tilt : str, optional + x_tilt : str | float, optional Tilt of the trajectory over the :math:`k_x`-axis, by default `None`. - y_tilt : str, optional + y_tilt : str | float, optional Tilt of the trajectory over the :math:`k_y`-axis, by default `None`. - z_tilt : str, optional + z_tilt : str | float, optional Tilt of the trajectory over the :math:`k_z`-axis, by default `None`. Returns ------- - array_like + NDArray Rotated trajectory. """ # Check dimensionality and initialize output @@ -96,33 +111,33 @@ def rotate(trajectory, nb_rotations, x_tilt=None, y_tilt=None, z_tilt=None): def precess( - trajectory, - nb_rotations, - tilt="golden", - half_sphere=False, - partition="axial", - axis=None, -): + trajectory: NDArray, + nb_rotations: int, + tilt: str | float = "golden", + half_sphere: bool = False, + partition: Literal["axial", "polar"] = "axial", + axis: int | NDArray | None = None, +) -> NDArray: """Rotate trajectories as a precession around the :math:`k_z`-axis. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory in 2D or 3D to rotate. nb_rotations : int Number of rotations repeating the provided trajectory while precessing. - tilt : str, optional + tilt : str | float, optional Angle tilt between consecutive rotations around the :math:`k_z`-axis, by default "golden". half_sphere : bool, optional Whether the precession should be limited to the upper half of the k-space sphere. It is typically used for in-out trajectories or planes. - partition : str, optional + partition : Literal["axial", "polar"], optional Partition type between an "axial" or "polar" split of the :math:`k_z`-axis, designating whether the axis should be fragmented by radius or angle respectively, by default "axial". - axis : int, array_like, optional + axis : int, NDArray, optional Axis selected for alignment reference when rotating the trajectory around the :math:`k_z`-axis, generally corresponding to the shot direction for single shot ``trajectory`` inputs. It can either @@ -132,7 +147,7 @@ def precess( Returns ------- - array_like + NDArray Precessed trajectory. """ # Check for partition option error @@ -175,22 +190,22 @@ def precess( def conify( - trajectory, - nb_cones, - z_tilt=None, - in_out=False, - max_angle=np.pi / 2, - borderless=True, -): + trajectory: NDArray, + nb_cones: int, + z_tilt: str | float | None = None, + in_out: bool = False, + max_angle: float = np.pi / 2, + borderless: bool = True, +) -> NDArray: """Distort 2D or 3D trajectories into cones along the :math:`k_z`-axis. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to conify. nb_cones : int Number of cones repeating the provided trajectory. - z_tilt : str, optional + z_tilt : str | float, optional Tilt of the trajectory over the :math:`k_z`-axis, by default `None`. in_out : bool, optional Whether to account for the in-out nature of some trajectories @@ -203,7 +218,7 @@ def conify( Returns ------- - array_like + NDArray Conified trajectory. """ # Check dimensionality and initialize output @@ -253,7 +268,13 @@ def conify( return new_trajectory -def epify(trajectory, Ns_transitions, nb_trains, reverse_odd_shots=False): +def epify( + trajectory: NDArray, + Ns_transitions: int, + nb_trains: int, + *, + reverse_odd_shots: bool = False, +) -> NDArray: """Create multi-readout shots from trajectory composed of single-readouts. Assemble multiple single-readout shots together by adding transition @@ -261,7 +282,7 @@ def epify(trajectory, Ns_transitions, nb_trains, reverse_odd_shots=False): Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to change by prolonging and merging the shots. Ns_transitions : int Number of samples/steps between the merged readouts. @@ -274,7 +295,7 @@ def epify(trajectory, Ns_transitions, nb_trains, reverse_odd_shots=False): Returns ------- - array_like + NDArray Trajectory with fewer but longer multi-readout shots. """ Nc, Ns, Nd = trajectory.shape @@ -302,12 +323,10 @@ def epify(trajectory, Ns_transitions, nb_trains, reverse_odd_shots=False): for i_c in range(nb_trains): spline = CubicSpline(source_sample_ids, np.concatenate(trajectory[i_c], axis=0)) assembled_trajectory.append(spline(target_sample_ids)) - assembled_trajectory = np.array(assembled_trajectory) - - return assembled_trajectory + return np.array(assembled_trajectory) -def unepify(trajectory, Ns_readouts, Ns_transitions): +def unepify(trajectory: NDArray, Ns_readouts: int, Ns_transitions: int) -> NDArray: """Recover single-readout shots from multi-readout trajectory. Reformat an EPI-like trajectory with multiple readouts and transitions @@ -319,7 +338,7 @@ def unepify(trajectory, Ns_readouts, Ns_transitions): Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to reduce by discarding transitions between readouts. Ns_readouts : int Number of samples within a single readout. @@ -328,7 +347,7 @@ def unepify(trajectory, Ns_readouts, Ns_transitions): Returns ------- - array_like + NDArray Trajectory with more but shorter single shots. """ Nc, Ns, Nd = trajectory.shape @@ -349,7 +368,7 @@ def unepify(trajectory, Ns_readouts, Ns_transitions): return trajectory -def prewind(trajectory, Ns_transitions): +def prewind(trajectory: NDArray, Ns_transitions: int) -> NDArray: """Add pre-winding/positioning to the trajectory. The trajectory is extended to start before the readout @@ -358,7 +377,7 @@ def prewind(trajectory, Ns_transitions): Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to extend with rewind gradients. Ns_transitions : int Number of pre-winding/positioning steps used to leave the @@ -366,7 +385,7 @@ def prewind(trajectory, Ns_transitions): Returns ------- - array_like + NDArray Extended trajectory with pre-winding/positioning. """ Nc, Ns, Nd = trajectory.shape @@ -384,12 +403,10 @@ def prewind(trajectory, Ns_transitions): np.concatenate([np.zeros((2, Nd)), trajectory[i_c]], axis=0), ) assembled_trajectory.append(spline(target_sample_ids)) - assembled_trajectory = np.array(assembled_trajectory) - - return assembled_trajectory + return np.array(assembled_trajectory) -def rewind(trajectory, Ns_transitions): +def rewind(trajectory: NDArray, Ns_transitions: int) -> NDArray: """Add rewinding to the trajectory. The trajectory is extended to come back to the k-space center @@ -397,14 +414,14 @@ def rewind(trajectory, Ns_transitions): Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to extend with rewind gradients. Ns_transitions : int Number of rewinding steps used to come back to the k-space center. Returns ------- - array_like + NDArray Extended trajectory with rewinding. """ Nc, Ns, Nd = trajectory.shape @@ -424,29 +441,31 @@ def rewind(trajectory, Ns_transitions): np.concatenate([trajectory[i_c], np.zeros((2, Nd))], axis=0), ) assembled_trajectory.append(spline(target_sample_ids)) - assembled_trajectory = np.array(assembled_trajectory) + return np.array(assembled_trajectory) - return assembled_trajectory - -def oversample(trajectory, new_Ns, kind="cubic"): +def oversample( + trajectory: NDArray, + new_Ns: int, + kind: Literal["linear", "quadratic", "cubic"] = "cubic", +) -> NDArray: """ Resample a trajectory to increase the number of samples using interpolation. Parameters ---------- - trajectory : np.ndarray + trajectory : NDArray The original trajectory array, where interpolation is applied along the second axis. new_Ns : int The desired number of samples in the resampled trajectory. - kind : str, optional + kind : Literal, optional The type of interpolation to use, such as 'linear', 'quadratic', or 'cubic', by default "cubic". Returns ------- - np.ndarray + NDArray The resampled trajectory array with ``new_Ns`` points along the second axis. Notes @@ -475,32 +494,36 @@ def oversample(trajectory, new_Ns, kind="cubic"): def stack_spherically( - trajectory_func, Nc, nb_stacks, z_tilt=None, hard_bounded=True, **traj_kwargs -): + trajectory_func: Callable[..., NDArray], + Nc: int, + nb_stacks: int, + z_tilt: str | float | None = None, + hard_bounded: bool = True, + **traj_kwargs: Any, # noqa ANN401 +) -> NDArray: """Stack 2D or 3D trajectories over the :math:`k_z`-axis to make a sphere. Parameters ---------- - trajectory_func : function + trajectory_func : Callable[..., NDArray] Trajectory function that should return an array-like with the usual (Nc, Ns, Nd) size. Nc : int - Number of shots to use for the whole spherically - stacked trajectory. + Number of shots to use for the whole spherically stacked trajectory. nb_stacks : int Number of stacks of trajectories. - z_tilt : str, optional + z_tilt : str | float, optional Tilt of the stacks, by default `None`. hard_bounded : bool, optional - Whether the stacks should be strictly within the limits of the k-space, - by default `True`. - **kwargs - Trajectory initialization parameters for the function provided - with `trajectory_func`. + Whether the stacks should be strictly within the limits + of the k-space, by default `True`. + **traj_kwargs : Any + Trajectory initialization parameters for the function + provided with `trajectory_func`. Returns ------- - array_like + NDArray Stacked trajectory. """ # Handle argument errors @@ -547,52 +570,45 @@ def stack_spherically( # Concatenate or handle varying Ns value Ns_values = np.array([stk.shape[1] for stk in new_trajectory]) if (Ns_values == Ns_values[0]).all(): - new_trajectory = np.concatenate(new_trajectory, axis=0) - new_trajectory = new_trajectory.reshape(Nc, Ns_values[0], 3) - else: - new_trajectory = np.concatenate( - [stk.reshape((-1, 3)) for stk in new_trajectory], axis=0 - ) - - return new_trajectory + output = np.concatenate(new_trajectory, axis=0) + return output.reshape(Nc, Ns_values[0], 3) + return np.concatenate([stk.reshape((-1, 3)) for stk in new_trajectory], axis=0) def shellify( - trajectory_func, - Nc, - nb_shells, - z_tilt="golden", - hemisphere_mode="symmetric", - **traj_kwargs, -): + trajectory_func: Callable[..., NDArray], + Nc: int, + nb_shells: int, + z_tilt: str | float = "golden", + hemisphere_mode: Literal["symmetric", "reversed"] = "symmetric", + **traj_kwargs: Any, # noqa ANN401 +) -> NDArray: """Stack 2D or 3D trajectories over the :math:`k_z`-axis to make a sphere. Parameters ---------- - trajectory_func : function - Trajectory function that should return an array-like - with the usual (Nc, Ns, Nd) size. + trajectory_func : Callable[..., NDArray] + Trajectory function that should return an array-like with the usual + (Nc, Ns, Nd) size. Nc : int - Number of shots to use for the whole spherically - stacked trajectory. + Number of shots to use for the whole spherically stacked trajectory. nb_shells : int - Number of shells of distorded trajectories. - z_tilt : str, float, optional + Number of shells of distorted trajectories. + z_tilt : str | float, optional Tilt of the shells, by default "golden". - hemisphere_mode : str, optional - Define how the lower hemisphere should be oriented - relatively to the upper one, with "symmetric" providing - a :math:`k_x-k_y` planar symmetry by changing the polar angle, - and with "reversed" promoting continuity (for example - in spirals) by reversing the azimuth angle. + hemisphere_mode : Literal["symmetric", "reversed"], optional + Define how the lower hemisphere should be oriented relatively to the + upper one, with "symmetric" providing a :math:`k_x-k_y` planar symmetry + by changing the polar angle, and with "reversed" promoting continuity + (for example in spirals) by reversing the azimuth angle. The default is "symmetric". - **kwargs - Trajectory initialization parameters for the function provided - with `trajectory_func`. + **traj_kwargs : Any + Trajectory initialization parameters for the function + provided with `trajectory_func`. Returns ------- - array_like + NDArray Concentric shell trajectory. """ # Handle argument errors @@ -643,14 +659,9 @@ def shellify( # Concatenate or handle varying Ns value Ns_values = np.array([hem.shape[1] for hem in new_trajectory]) if (Ns_values == Ns_values[0]).all(): - new_trajectory = np.concatenate(new_trajectory, axis=0) - new_trajectory = new_trajectory.reshape(Nc, Ns_values[0], 3) - else: - new_trajectory = np.concatenate( - [hem.reshape((-1, 3)) for hem in new_trajectory], axis=0 - ) - - return new_trajectory + output = np.concatenate(new_trajectory, axis=0) + return output.reshape(Nc, Ns_values[0], 3) + return np.concatenate([hem.reshape((-1, 3)) for hem in new_trajectory], axis=0) ######### @@ -658,7 +669,9 @@ def shellify( ######### -def duplicate_along_axes(trajectory, axes=(0, 1, 2)): +def duplicate_along_axes( + trajectory: NDArray, axes: tuple[int, ...] = (0, 1, 2) +) -> NDArray: """ Duplicate a trajectory along the specified axes. @@ -668,14 +681,14 @@ def duplicate_along_axes(trajectory, axes=(0, 1, 2)): Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to duplicate. - axes : tuple, optional + axes : tuple[int, ...], optional Axes along which to duplicate the trajectory, by default (0, 1, 2) Returns ------- - array_like + NDArray Duplicated trajectory along the specified axes. """ # Copy input trajectory along other axes @@ -690,12 +703,24 @@ def duplicate_along_axes(trajectory, axes=(0, 1, 2)): dp_trajectory = np.copy(trajectory) dp_trajectory[..., [2, 0]] = dp_trajectory[..., [0, 2]] new_trajectory.append(dp_trajectory) - new_trajectory = np.concatenate(new_trajectory, axis=0) - return new_trajectory + return np.concatenate(new_trajectory, axis=0) -def _radialize_center_out(trajectory, nb_samples): - """Radialize a trajectory from the center to the outside.""" +def _radialize_center_out(trajectory: NDArray, nb_samples: int) -> NDArray: + """Radialize a trajectory from the center to the outside. + + Parameters + ---------- + trajectory : NDArray + Trajectory to radialize. + nb_samples : int + Number of samples to radialize from the center. + + Returns + ------- + NDArray + Radialized trajectory. + """ Nc, Ns = trajectory.shape[:2] new_trajectory = np.copy(trajectory) for i in range(Nc): @@ -706,8 +731,21 @@ def _radialize_center_out(trajectory, nb_samples): return new_trajectory -def _radialize_in_out(trajectory, nb_samples): - """Radialize a trajectory from the inside to the outside.""" +def _radialize_in_out(trajectory: NDArray, nb_samples: int) -> NDArray: + """Radialize a trajectory from the inside to the outside. + + Parameters + ---------- + trajectory : NDArray + Trajectory to radialize. + nb_samples : int + Number of samples to radialize from the inside out. + + Returns + ------- + NDArray + Radialized trajectory. + """ Nc, Ns = trajectory.shape[:2] new_trajectory = np.copy(trajectory) first, half, second = (Ns - nb_samples) // 2, Ns // 2, (Ns + nb_samples) // 2 @@ -723,20 +761,26 @@ def _radialize_in_out(trajectory, nb_samples): return new_trajectory -def radialize_center(trajectory, nb_samples, in_out=False): +def radialize_center( + trajectory: NDArray, nb_samples: int, in_out: bool = False +) -> NDArray: """Radialize a trajectory. Parameters ---------- - trajectory : array_like + trajectory : NDArray Trajectory to radialize. nb_samples : int Number of samples to keep. in_out : bool, optional Whether the radialization is from the inside to the outside, by default False + + Returns + ------- + NDArray + Radialized trajectory. """ # Make nb_samples into straight lines around the center if in_out: return _radialize_in_out(trajectory, nb_samples) - else: - return _radialize_center_out(trajectory, nb_samples) + return _radialize_center_out(trajectory, nb_samples) diff --git a/src/mrinufft/trajectories/trajectory2D.py b/src/mrinufft/trajectories/trajectory2D.py index e930e90b3..a9029e25c 100644 --- a/src/mrinufft/trajectories/trajectory2D.py +++ b/src/mrinufft/trajectories/trajectory2D.py @@ -1,7 +1,10 @@ """Functions to initialize 2D trajectories.""" +from typing import Any, Literal + import numpy as np import numpy.linalg as nl +from numpy.typing import NDArray from scipy.interpolate import CubicSpline from .gradients import patch_center_anomaly @@ -14,7 +17,9 @@ ##################### -def initialize_2D_radial(Nc, Ns, tilt="uniform", in_out=False): +def initialize_2D_radial( + Nc: int, Ns: int, tilt: str | float = "uniform", in_out: bool = False +) -> NDArray: """Initialize a 2D radial trajectory. Parameters @@ -23,14 +28,14 @@ def initialize_2D_radial(Nc, Ns, tilt="uniform", in_out=False): Number of shots Ns : int Number of samples per shot - tilt : str, float, optional + tilt : str | float, optional Tilt of the shots, by default "uniform" in_out : bool, optional Whether to start from the center or not, by default False Returns ------- - array_like + NDArray 2D radial trajectory """ # Initialize a first shot @@ -47,14 +52,14 @@ def initialize_2D_radial(Nc, Ns, tilt="uniform", in_out=False): def initialize_2D_spiral( - Nc, - Ns, - tilt="uniform", - in_out=False, - nb_revolutions=1, - spiral="archimedes", - patch_center=True, -): + Nc: int, + Ns: int, + tilt: str | float = "uniform", + in_out: bool = False, + nb_revolutions: float = 1.0, + spiral: str | float = "archimedes", + patch_center: bool = True, +) -> NDArray: """Initialize a 2D algebraic spiral trajectory. A generalized function that generates algebraic spirals defined @@ -68,13 +73,13 @@ def initialize_2D_spiral( Number of shots Ns : int Number of samples per shot - tilt : str, float, optional + tilt : Literal, float, optional Tilt of the shots, by default "uniform" in_out : bool, optional Whether to start from the center or not, by default False - nb_revolutions : int, optional + nb_revolutions : float, optional Number of revolutions, by default 1 - spiral : str, float, optional + spiral : Literal, float, optional Spiral type or algebraic power, by default "archimedes" patch_center : bool, optional Whether the spiral anomaly at the center should be patched @@ -82,7 +87,7 @@ def initialize_2D_spiral( Returns ------- - array_like + NDArray 2D spiral trajectory Raises @@ -109,11 +114,18 @@ def initialize_2D_spiral( # Algebraic spirals with power coefficients superior to 1 # have a non-monotonic gradient norm when varying the angle # over [0, +inf) - def _update_shot(angles, radius, *args): + def _update_shot( + angles: NDArray, radius: NDArray, *args: Any # noqa ANN401 + ) -> NDArray: shot = np.sign(angles) * np.abs(radius) * np.exp(1j * np.abs(angles)) return np.stack([shot.real, shot.imag], axis=-1) - def _update_parameters(single_shot, angles, radius, spiral_power): + def _update_parameters( + single_shot: NDArray, + angles: NDArray, + radius: NDArray, + spiral_power: float, + ) -> tuple[NDArray, NDArray, float]: radius = nl.norm(single_shot, axis=-1) angles = np.sign(angles) * np.abs(radius) ** (1 / spiral_power) return angles, radius, spiral_power @@ -144,7 +156,9 @@ def _update_parameters(single_shot, angles, radius, spiral_power): return trajectory -def initialize_2D_fibonacci_spiral(Nc, Ns, spiral_reduction=1, patch_center=True): +def initialize_2D_fibonacci_spiral( + Nc: int, Ns: int, spiral_reduction: float = 1, patch_center: bool = True +) -> NDArray: """Initialize a 2D Fibonacci spiral trajectory. A non-algebraic spiral trajectory based on the Fibonacci sequence, @@ -168,7 +182,7 @@ def initialize_2D_fibonacci_spiral(Nc, Ns, spiral_reduction=1, patch_center=True Returns ------- - array_like + NDArray 2D Fibonacci spiral trajectory References @@ -213,7 +227,14 @@ def initialize_2D_fibonacci_spiral(Nc, Ns, spiral_reduction=1, patch_center=True return trajectory -def initialize_2D_cones(Nc, Ns, tilt="uniform", in_out=False, nb_zigzags=5, width=1): +def initialize_2D_cones( + Nc: int, + Ns: int, + tilt: str | float = "uniform", + in_out: bool = False, + nb_zigzags: float = 5, + width: float = 1, +) -> NDArray: """Initialize a 2D cone trajectory. Parameters @@ -222,7 +243,7 @@ def initialize_2D_cones(Nc, Ns, tilt="uniform", in_out=False, nb_zigzags=5, widt Number of shots Ns : int Number of samples per shot - tilt : str, optional + tilt : str | float, optional Tilt of the shots, by default "uniform" in_out : bool, optional Whether to start from the center or not, by default False @@ -233,7 +254,7 @@ def initialize_2D_cones(Nc, Ns, tilt="uniform", in_out=False, nb_zigzags=5, widt Returns ------- - array_like + NDArray 2D cone trajectory """ @@ -253,8 +274,13 @@ def initialize_2D_cones(Nc, Ns, tilt="uniform", in_out=False, nb_zigzags=5, widt def initialize_2D_sinusoide( - Nc, Ns, tilt="uniform", in_out=False, nb_zigzags=5, width=1 -): + Nc: int, + Ns: int, + tilt: str | float = "uniform", + in_out: bool = False, + nb_zigzags: float = 5, + width: float = 1, +) -> NDArray: """Initialize a 2D sinusoide trajectory. Parameters @@ -263,7 +289,7 @@ def initialize_2D_sinusoide( Number of shots Ns : int Number of samples per shot - tilt : str, float, optional + tilt : str | float, optional Tilt of the shots, by default "uniform" in_out : bool, optional Whether to start from the center or not, by default False @@ -274,7 +300,7 @@ def initialize_2D_sinusoide( Returns ------- - array_like + NDArray 2D sinusoide trajectory """ @@ -293,7 +319,7 @@ def initialize_2D_sinusoide( return trajectory -def initialize_2D_propeller(Nc, Ns, nb_strips): +def initialize_2D_propeller(Nc: int, Ns: int, nb_strips: int) -> NDArray: """Initialize a 2D PROPELLER trajectory, as proposed in [Pip99]_. The PROPELLER trajectory is generally used along a specific @@ -341,7 +367,7 @@ def initialize_2D_propeller(Nc, Ns, nb_strips): return KMAX * trajectory -def initialize_2D_rings(Nc, Ns, nb_rings): +def initialize_2D_rings(Nc: int, Ns: int, nb_rings: int) -> NDArray: """Initialize a 2D ring trajectory, as proposed in [HHN08]_. Parameters @@ -355,7 +381,7 @@ def initialize_2D_rings(Nc, Ns, nb_rings): Returns ------- - array_like + NDArray 2D ring trajectory References @@ -387,7 +413,9 @@ def initialize_2D_rings(Nc, Ns, nb_rings): return KMAX * np.array(trajectory) -def initialize_2D_rosette(Nc, Ns, in_out=False, coprime_index=0): +def initialize_2D_rosette( + Nc: int, Ns: int, in_out: bool = False, coprime_index: int = 0 +) -> NDArray: """Initialize a 2D rosette trajectory. Parameters @@ -403,7 +431,7 @@ def initialize_2D_rosette(Nc, Ns, in_out=False, coprime_index=0): Returns ------- - array_like + NDArray 2D rosette trajectory """ @@ -428,7 +456,9 @@ def initialize_2D_rosette(Nc, Ns, in_out=False, coprime_index=0): return trajectory -def initialize_2D_polar_lissajous(Nc, Ns, in_out=False, nb_segments=1, coprime_index=0): +def initialize_2D_polar_lissajous( + Nc: int, Ns: int, in_out: bool = False, nb_segments: int = 1, coprime_index: int = 0 +) -> NDArray: """Initialize a 2D polar Lissajous trajectory. Parameters @@ -446,7 +476,7 @@ def initialize_2D_polar_lissajous(Nc, Ns, in_out=False, nb_segments=1, coprime_i Returns ------- - array_like + NDArray 2D polar Lissajous trajectory """ # Adapt the parameters to subcases @@ -482,7 +512,7 @@ def initialize_2D_polar_lissajous(Nc, Ns, in_out=False, nb_segments=1, coprime_i ######################### -def initialize_2D_lissajous(Nc, Ns, density=1): +def initialize_2D_lissajous(Nc: int, Ns: int, density: float = 1) -> NDArray: """Initialize a 2D Lissajous trajectory. Parameters @@ -496,7 +526,7 @@ def initialize_2D_lissajous(Nc, Ns, density=1): Returns ------- - array_like + NDArray 2D Lissajous trajectory """ # Define the whole curve in Cartesian coordinates @@ -512,7 +542,9 @@ def initialize_2D_lissajous(Nc, Ns, density=1): return trajectory -def initialize_2D_waves(Nc, Ns, nb_zigzags=5, width=1): +def initialize_2D_waves( + Nc: int, Ns: int, nb_zigzags: float = 5, width: float = 1 +) -> NDArray: """Initialize a 2D waves trajectory. Parameters @@ -528,7 +560,7 @@ def initialize_2D_waves(Nc, Ns, nb_zigzags=5, width=1): Returns ------- - array_like + NDArray 2D waves trajectory """ # Initialize a first shot diff --git a/src/mrinufft/trajectories/trajectory3D.py b/src/mrinufft/trajectories/trajectory3D.py index cdcf4a0ad..5ec9f2bca 100644 --- a/src/mrinufft/trajectories/trajectory3D.py +++ b/src/mrinufft/trajectories/trajectory3D.py @@ -1,9 +1,11 @@ """Functions to initialize 3D trajectories.""" from functools import partial +from typing import Literal import numpy as np import numpy.linalg as nl +from numpy.typing import NDArray from scipy.special import ellipj, ellipk from .maths import ( @@ -17,14 +19,16 @@ ) from .tools import conify, duplicate_along_axes, epify, precess, stack from .trajectory2D import initialize_2D_radial, initialize_2D_spiral -from .utils import KMAX, Packings, initialize_shape_norm, initialize_tilt +from .utils import KMAX, Packings, Spirals, initialize_shape_norm, initialize_tilt ############## # 3D RADIALS # ############## -def initialize_3D_phyllotaxis_radial(Nc, Ns, nb_interleaves=1, in_out=False): +def initialize_3D_phyllotaxis_radial( + Nc: int, Ns: int, nb_interleaves: int = 1, in_out: bool = False +) -> NDArray: """Initialize 3D radial trajectories with phyllotactic structure. The radial shots are oriented according to a Fibonacci sphere @@ -53,7 +57,7 @@ def initialize_3D_phyllotaxis_radial(Nc, Ns, nb_interleaves=1, in_out=False): Returns ------- - array_like + NDArray 3D phyllotaxis radial trajectory References @@ -71,7 +75,9 @@ def initialize_3D_phyllotaxis_radial(Nc, Ns, nb_interleaves=1, in_out=False): return trajectory -def initialize_3D_golden_means_radial(Nc, Ns, in_out=False): +def initialize_3D_golden_means_radial( + Nc: int, Ns: int, in_out: bool = False +) -> NDArray: """Initialize 3D radial trajectories with golden means-based structure. The radial shots are oriented using multidimensional golden means, @@ -95,7 +101,7 @@ def initialize_3D_golden_means_radial(Nc, Ns, in_out=False): Returns ------- - array_like + NDArray 3D golden means radial trajectory References @@ -123,7 +129,9 @@ def initialize_3D_golden_means_radial(Nc, Ns, in_out=False): return KMAX * trajectory -def initialize_3D_wong_radial(Nc, Ns, nb_interleaves=1, in_out=False): +def initialize_3D_wong_radial( + Nc: int, Ns: int, nb_interleaves: int = 1, in_out: bool = False +) -> NDArray: """Initialize 3D radial trajectories with a spiral structure. The radial shots are oriented according to an archimedean spiral @@ -149,7 +157,7 @@ def initialize_3D_wong_radial(Nc, Ns, nb_interleaves=1, in_out=False): Returns ------- - array_like + NDArray 3D Wong radial trajectory References @@ -181,7 +189,9 @@ def initialize_3D_wong_radial(Nc, Ns, nb_interleaves=1, in_out=False): return trajectory -def initialize_3D_park_radial(Nc, Ns, nb_interleaves=1, in_out=False): +def initialize_3D_park_radial( + Nc: int, Ns: int, nb_interleaves: int = 1, in_out: bool = False +) -> NDArray: """Initialize 3D radial trajectories with a spiral structure. The radial shots are oriented according to an archimedean spiral @@ -208,7 +218,7 @@ def initialize_3D_park_radial(Nc, Ns, nb_interleaves=1, in_out=False): Returns ------- - array_like + NDArray 3D Park radial trajectory References @@ -233,8 +243,14 @@ def initialize_3D_park_radial(Nc, Ns, nb_interleaves=1, in_out=False): def initialize_3D_cones( - Nc, Ns, tilt="golden", in_out=False, nb_zigzags=5, spiral="archimedes", width=1 -): + Nc: int, + Ns: int, + tilt: str | float = "golden", + in_out: bool = False, + nb_zigzags: float = 5, + spiral: str | float = "archimedes", + width: float = 1, +) -> NDArray: """Initialize 3D trajectories with cones. Initialize a trajectory consisting of 3D cones duplicated @@ -264,7 +280,7 @@ def initialize_3D_cones( Returns ------- - array_like + NDArray 3D cones trajectory References @@ -274,7 +290,7 @@ def initialize_3D_cones( Journal of mathematical chemistry 6, no. 1 (1991): 325-349. """ # Initialize first spiral - spiral = initialize_2D_spiral( + single_spiral = initialize_2D_spiral( Nc=1, Ns=Ns, spiral=spiral, @@ -290,8 +306,8 @@ def initialize_3D_cones( # Initialize first cone ## Create three cones for proper partitioning, but only one is needed - cone = conify( - spiral, + cones = conify( + single_spiral, nb_cones=3, z_tilt=None, in_out=in_out, @@ -301,22 +317,27 @@ def initialize_3D_cones( # Apply precession to the first cone trajectory = precess( - cone, tilt=tilt, nb_rotations=Nc, half_sphere=in_out, partition="axial", axis=2 + cones, + tilt=tilt, + nb_rotations=Nc, + half_sphere=in_out, + partition="axial", + axis=2, ) return trajectory def initialize_3D_floret( - Nc, - Ns, - in_out=False, - nb_revolutions=1, - spiral="fermat", - cone_tilt="golden", - max_angle=np.pi / 2, - axes=(2,), -): + Nc: int, + Ns: int, + in_out: bool = False, + nb_revolutions: float = 1, + spiral: str | float = "fermat", + cone_tilt: str | float = "golden", + max_angle: float = np.pi / 2, + axes: tuple[int, ...] = (2,), +) -> NDArray: """Initialize 3D trajectories with FLORET. This implementation is based on the work from [Pip+11]_. @@ -346,7 +367,7 @@ def initialize_3D_floret( Returns ------- - array_like + NDArray 3D FLORET trajectory References @@ -363,7 +384,7 @@ def initialize_3D_floret( raise ValueError("Nc should be divisible by len(axes).") # Initialize first spiral - spiral = initialize_2D_spiral( + single_spiral = initialize_2D_spiral( Nc=1, Ns=Ns, spiral=spiral, @@ -372,8 +393,8 @@ def initialize_3D_floret( ) # Initialize first cone - cone = conify( - spiral, + cones = conify( + single_spiral, nb_cones=Nc_per_axis, z_tilt=cone_tilt, in_out=in_out, @@ -381,20 +402,20 @@ def initialize_3D_floret( ) # Duplicate cone along axes - axes = [2 - ax for ax in axes] # Default axis is kz, not kx - trajectory = duplicate_along_axes(cone, axes=axes) + axes = tuple(2 - ax for ax in axes) # Default axis is kz, not kx + trajectory = duplicate_along_axes(cones, axes=axes) return trajectory def initialize_3D_wave_caipi( - Nc, - Ns, - nb_revolutions=5, - width=1, - packing="triangular", - shape="square", - spacing=(1, 1), -): + Nc: int, + Ns: int, + nb_revolutions: float = 5, + width: float = 1, + packing: str = "triangular", + shape: str | float = "square", + spacing: tuple[int, int] = (1, 1), +) -> NDArray: """Initialize 3D trajectories with Wave-CAIPI. This implementation is based on the work from [Bil+15]_. @@ -415,20 +436,20 @@ def initialize_3D_wave_caipi( Packing method used to position the helices: "triangular"/"hexagonal", "square", "circular" or "random"/"uniform", by default "triangular". - shape : str or float, optional + shape : str | float, optional Shape over the 2D :math:`k_x-k_y` plane to pack with shots, either defined as `str` ("circle", "square", "diamond") or as `float` through p-norms following the conventions of the `ord` parameter from `numpy.linalg.norm`, by default "circle". - spacing : tuple(int, int) + spacing : tuple[int, int] Spacing between helices over the 2D :math:`k_x-k_y` plane normalized similarly to `width` to correspond to helix diameters, by default (1, 1). Returns ------- - array_like + NDArray 3D wave-CAIPI trajectory References @@ -439,7 +460,6 @@ def initialize_3D_wave_caipi( Magnetic resonance in medicine 73, no. 6 (2015): 2152-2162. """ trajectory = np.zeros((Nc, Ns, 3)) - spacing = np.array(spacing) # Initialize first shot angles = nb_revolutions * 2 * np.pi * np.arange(0, Ns) / Ns @@ -448,11 +468,11 @@ def initialize_3D_wave_caipi( trajectory[0, :, 2] = np.linspace(-1, 1, Ns) # Choose the helix positions according to packing - packing = Packings[packing] + packing_enum = Packings[packing] side = 2 * int(np.ceil(np.sqrt(Nc))) * np.max(spacing) - if packing == Packings.RANDOM: + if packing_enum == Packings.RANDOM: positions = 2 * side * (np.random.random((side * side, 2)) - 0.5) - elif packing == Packings.CIRCLE: + elif packing_enum == Packings.CIRCLE: positions = [[0, 0]] counter = 0 while len(positions) < side**2: @@ -466,21 +486,21 @@ def initialize_3D_wave_caipi( positions = np.concatenate( [positions, np.array([circle.real, circle.imag]).T], axis=0 ) - elif packing in [Packings.SQUARE, Packings.TRIANGLE, Packings.HEXAGONAL]: + elif packing_enum in [Packings.SQUARE, Packings.TRIANGLE, Packings.HEXAGONAL]: # Square packing or initialize hexagonal/triangular packing px, py = np.meshgrid( np.arange(-side + 1, side, 2), np.arange(-side + 1, side, 2) ) positions = np.stack([px.flatten(), py.flatten()], axis=-1).astype(float) - if packing in [Packings.HEXAGON, Packings.TRIANGLE]: + if packing_enum in [Packings.HEXAGON, Packings.TRIANGLE]: # Hexagonal/triangular packing based on square packing positions[::2, 1] += 1 / 2 positions[1::2, 1] -= 1 / 2 ratio = nl.norm(np.diff(positions[:2], axis=-1)) positions[:, 0] /= ratio / 2 - if packing == Packings.FIBONACCI: + if packing_enum == Packings.FIBONACCI: # Estimate helix width based on the k-space 2D surface # and an optimal circle packing positions = np.sqrt( @@ -490,7 +510,7 @@ def initialize_3D_wave_caipi( # Remove points by distance to fit both shape and Nc main_order = initialize_shape_norm(shape) tie_order = 2 if (main_order != 2) else np.inf # breaking ties - positions = np.array(positions) * spacing + positions = np.array(positions) * np.array(spacing) positions = sorted(positions, key=partial(nl.norm, ord=tie_order)) positions = sorted(positions, key=partial(nl.norm, ord=main_order)) positions = positions[:Nc] @@ -506,14 +526,14 @@ def initialize_3D_wave_caipi( def initialize_3D_seiffert_spiral( - Nc, - Ns, - curve_index=0.2, - nb_revolutions=1, - axis_tilt="golden", - spiral_tilt="golden", - in_out=False, -): + Nc: int, + Ns: int, + curve_index: float = 0.2, + nb_revolutions: float = 1, + axis_tilt: str | float = "golden", + spiral_tilt: str | float = "golden", + in_out: bool = False, +) -> NDArray: """Initialize 3D trajectories with modulated Seiffert spirals. Initially introduced in [SMR18]_, but also proposed later as "Yarnball" @@ -543,7 +563,7 @@ def initialize_3D_seiffert_spiral( Returns ------- - array_like + NDArray 3D Seiffert spiral trajectory References @@ -611,8 +631,13 @@ def initialize_3D_seiffert_spiral( def initialize_3D_helical_shells( - Nc, Ns, nb_shells, spiral_reduction=1, shell_tilt="intergaps", shot_tilt="uniform" -): + Nc: int, + Ns: int, + nb_shells: int, + spiral_reduction: float = 1, + shell_tilt: str = "intergaps", + shot_tilt: str = "uniform", +) -> NDArray: """Initialize 3D trajectories with helical shells. The implementation follows the proposition from [YRB06]_ @@ -635,7 +660,7 @@ def initialize_3D_helical_shells( Returns ------- - array_like + NDArray 3D helical shell trajectory References @@ -694,8 +719,12 @@ def initialize_3D_helical_shells( def initialize_3D_annular_shells( - Nc, Ns, nb_shells, shell_tilt=np.pi, ring_tilt=np.pi / 2 -): + Nc: int, + Ns: int, + nb_shells: int, + shell_tilt: float = np.pi, + ring_tilt: float = np.pi / 2, +) -> NDArray: """Initialize 3D trajectories with annular shells. An exclusive trajectory inspired from the work proposed in [HM11]_. @@ -715,7 +744,7 @@ def initialize_3D_annular_shells( Returns ------- - array_like + NDArray 3D annular shell trajectory References @@ -807,14 +836,14 @@ def initialize_3D_annular_shells( def initialize_3D_seiffert_shells( - Nc, - Ns, - nb_shells, - curve_index=0.5, - nb_revolutions=1, - shell_tilt="uniform", - shot_tilt="uniform", -): + Nc: int, + Ns: int, + nb_shells: int, + curve_index: float = 0.5, + nb_revolutions: float = 1, + shell_tilt: str = "uniform", + shot_tilt: str = "uniform", +) -> NDArray: """Initialize 3D trajectories with Seiffert shells. The implementation is based on work from [Er00]_ and [Br09]_, @@ -841,7 +870,7 @@ def initialize_3D_seiffert_shells( Returns ------- - array_like + NDArray 3D Seiffert shell trajectory References @@ -905,15 +934,15 @@ def initialize_3D_seiffert_shells( def initialize_3D_turbine( - Nc, - Ns_readouts, - Ns_transitions, - nb_blades, - blade_tilt="uniform", - nb_trains="auto", - skip_factor=1, - in_out=True, -): + Nc: int, + Ns_readouts: int, + Ns_transitions: int, + nb_blades: int, + blade_tilt: str = "uniform", + nb_trains: int | Literal["auto"] = "auto", + skip_factor: int = 1, + in_out: bool = True, +) -> NDArray: """Initialize 3D TURBINE trajectory. This is an implementation of the TURBINE (Trajectory Using Radially @@ -938,7 +967,7 @@ def initialize_3D_turbine( Number of line stacks over the :math:`k_z`-axis axis blade_tilt : str, float, optional Tilt between individual blades, by default "uniform" - nb_trains : int, str, optional + nb_trains : int, Literal["auto"], optional Number of resulting shots, or readout trains, such that each of them will be composed of :math:`n` readouts with ``Nc = n * nb_trains``. If ``"auto"`` then ``nb_trains`` is set @@ -951,7 +980,7 @@ def initialize_3D_turbine( Returns ------- - array_like + NDArray 3D TURBINE trajectory References @@ -1013,17 +1042,17 @@ def initialize_3D_turbine( def initialize_3D_repi( - Nc, - Ns_readouts, - Ns_transitions, - nb_blades, - nb_blade_revolutions=0, - blade_tilt="uniform", - nb_spiral_revolutions=0, - spiral="archimedes", - nb_trains="auto", - in_out=True, -): + Nc: int, + Ns_readouts: int, + Ns_transitions: int, + nb_blades: int, + nb_blade_revolutions: float = 0, + blade_tilt: str = "uniform", + nb_spiral_revolutions: float = 0, + spiral: str = "archimedes", + nb_trains: int | Literal["auto"] = "auto", + in_out: bool = True, +) -> NDArray: """Initialize 3D REPI trajectory. This is an implementation of the REPI (Radial Echo Planar Imaging) @@ -1058,7 +1087,7 @@ def initialize_3D_repi( Number of revolutions of the spirals over the readouts, by default 0 spiral : str, float, optional Spiral type, by default "archimedes" - nb_trains : int, str + nb_trains : int, Literal["auto"], optional Number of trains dividing the readouts, such that each shot will be composed of `n` readouts with `Nc = n * nb_trains`. If "auto" then `nb_trains` is set to `nb_blades`. @@ -1067,7 +1096,7 @@ def initialize_3D_repi( Returns ------- - array_like + NDArray 3D REPI trajectory References diff --git a/src/mrinufft/trajectories/utils.py b/src/mrinufft/trajectories/utils.py index e2c9ff4fb..c1124e116 100644 --- a/src/mrinufft/trajectories/utils.py +++ b/src/mrinufft/trajectories/utils.py @@ -2,8 +2,10 @@ from enum import Enum, EnumMeta from numbers import Real +from typing import Any, Literal import numpy as np +from numpy.typing import NDArray ############# # CONSTANTS # @@ -26,11 +28,11 @@ class CaseInsensitiveEnumMeta(EnumMeta): """A case-insensitive EnumMeta.""" - def __getitem__(self, name): + def __getitem__(self, name: str) -> Enum: """Allow ``MyEnum['Member'] == MyEnum['MEMBER']`` .""" return super().__getitem__(name.upper()) - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: # noqa ANN401 """Allow ``MyEnum.Member == MyEnum.MEMBER`` .""" return super().__getattr__(name.upper()) @@ -121,7 +123,7 @@ class Tilts(str, Enum): class Packings(str, Enum, metaclass=CaseInsensitiveEnumMeta): """Enumerate available packing method for shots. - It is mostly use for wave-CAIPI trajectory + It is mostly used for wave-CAIPI trajectory See Also -------- @@ -150,15 +152,15 @@ class Packings(str, Enum, metaclass=CaseInsensitiveEnumMeta): def normalize_trajectory( - trajectory, - norm_factor=KMAX, - resolution=DEFAULT_RESOLUTION, -): + trajectory: NDArray, + norm_factor: float = KMAX, + resolution: float | NDArray = DEFAULT_RESOLUTION, +) -> NDArray: """Normalize an un-normalized/natural trajectory for NUFFT use. Parameters ---------- - trajectory : np.ndarray + trajectory : NDArray Un-normalized trajectory consisting of k-space coordinates in 2D or 3D. norm_factor : float, optional Trajectory normalization factor, by default KMAX. @@ -169,22 +171,22 @@ def normalize_trajectory( Returns ------- - trajectory : np.ndarray + trajectory : NDArray Normalized trajectory corresponding to `trajectory` input. """ return trajectory * norm_factor * (2 * resolution) def unnormalize_trajectory( - trajectory, - norm_factor=KMAX, - resolution=DEFAULT_RESOLUTION, -): + trajectory: NDArray, + norm_factor: float = KMAX, + resolution: float | NDArray = DEFAULT_RESOLUTION, +) -> NDArray: """Un-normalize a NUFFT-normalized trajectory. Parameters ---------- - trajectory : np.ndarray + trajectory : NDArray Normalized trajectory consisting of k-space coordinates in 2D or 3D. norm_factor : float, optional Trajectory normalization factor, by default KMAX. @@ -195,25 +197,25 @@ def unnormalize_trajectory( Returns ------- - trajectory : np.ndarray + trajectory : NDArray Un-normalized trajectory corresponding to `trajectory` input. """ return trajectory / norm_factor / (2 * resolution) def convert_trajectory_to_gradients( - trajectory, - norm_factor=KMAX, - resolution=DEFAULT_RESOLUTION, - raster_time=DEFAULT_RASTER_TIME, - gamma=Gammas.HYDROGEN, - get_final_positions=False, -): + trajectory: NDArray, + norm_factor: float = KMAX, + resolution: float | NDArray = DEFAULT_RESOLUTION, + raster_time: float = DEFAULT_RASTER_TIME, + gamma: float = Gammas.HYDROGEN, + get_final_positions: bool = False, +) -> tuple[NDArray, ...]: """Derive a normalized trajectory over time to provide gradients. Parameters ---------- - trajectory : np.ndarray + trajectory : NDArray Normalized trajectory consisting of k-space coordinates in 2D or 3D. norm_factor : float, optional Trajectory normalization factor, by default KMAX. @@ -234,7 +236,7 @@ def convert_trajectory_to_gradients( Returns ------- - gradients : np.ndarray + gradients : NDArray Gradients corresponding to `trajectory`. """ # Un-normalize the trajectory from NUFFT usage @@ -249,20 +251,20 @@ def convert_trajectory_to_gradients( def convert_gradients_to_trajectory( - gradients, - initial_positions=None, - norm_factor=KMAX, - resolution=DEFAULT_RESOLUTION, - raster_time=DEFAULT_RASTER_TIME, - gamma=Gammas.HYDROGEN, -): + gradients: NDArray, + initial_positions: NDArray | None = None, + norm_factor: float = KMAX, + resolution: float | NDArray = DEFAULT_RESOLUTION, + raster_time: float = DEFAULT_RASTER_TIME, + gamma: float = Gammas.HYDROGEN, +) -> NDArray: """Integrate gradients over time to provide a normalized trajectory. Parameters ---------- - gradients : np.ndarray + gradients : NDArray Gradients over 2 or 3 directions. - initial_positions: np.ndarray, optional + initial_positions: NDArray, optional Positions in k-space at the beginning of the readout window. The default is `None`. norm_factor : float, optional @@ -281,7 +283,7 @@ def convert_gradients_to_trajectory( Returns ------- - trajectory : np.ndarray + trajectory : NDArray Normalized trajectory corresponding to `gradients`. """ # Handle no initial positions @@ -299,14 +301,14 @@ def convert_gradients_to_trajectory( def convert_gradients_to_slew_rates( - gradients, - raster_time=DEFAULT_RASTER_TIME, -): + gradients: NDArray, + raster_time: float = DEFAULT_RASTER_TIME, +) -> tuple[NDArray, NDArray]: """Derive the gradients over time to provide slew rates. Parameters ---------- - gradients : np.ndarray + gradients : NDArray Gradients over 2 or 3 directions. raster_time : float, optional Amount of time between the acquisition of two @@ -315,9 +317,9 @@ def convert_gradients_to_slew_rates( Returns ------- - slewrates : np.ndarray + slewrates : NDArray Slew rates corresponding to `gradients`. - initial_gradients : np.ndarray + initial_gradients : NDArray Gradients at the beginning of the readout window. """ # Compute slew rates and starting gradients @@ -327,17 +329,17 @@ def convert_gradients_to_slew_rates( def convert_slew_rates_to_gradients( - slewrates, - initial_gradients=None, - raster_time=DEFAULT_RASTER_TIME, -): + slewrates: NDArray, + initial_gradients: NDArray | None = None, + raster_time: float = DEFAULT_RASTER_TIME, +) -> NDArray: """Integrate slew rates over time to provide gradients. Parameters ---------- - slewrates : np.ndarray + slewrates : NDArray Slew rates over 2 or 3 directions. - initial_gradients: np.ndarray, optional + initial_gradients: NDArray, optional Gradients at the beginning of the readout window. The default is `None`. raster_time : float, optional @@ -347,7 +349,7 @@ def convert_slew_rates_to_gradients( Returns ------- - gradients : np.ndarray + gradients : NDArray Gradients corresponding to `slewrates`. """ # Handle no initial gradients @@ -362,17 +364,17 @@ def convert_slew_rates_to_gradients( def compute_gradients_and_slew_rates( - trajectory, - norm_factor=KMAX, - resolution=DEFAULT_RESOLUTION, - raster_time=DEFAULT_RASTER_TIME, - gamma=Gammas.HYDROGEN, -): + trajectory: NDArray, + norm_factor: float = KMAX, + resolution: float | NDArray = DEFAULT_RESOLUTION, + raster_time: float = DEFAULT_RASTER_TIME, + gamma: float = Gammas.HYDROGEN, +) -> tuple[NDArray, NDArray]: """Compute the gradients and slew rates from a normalized trajectory. Parameters ---------- - trajectory : np.ndarray + trajectory : NDArray Normalized trajectory consisting of k-space coordinates in 2D or 3D. norm_factor : float, optional Trajectory normalization factor, by default KMAX. @@ -390,9 +392,9 @@ def compute_gradients_and_slew_rates( Returns ------- - gradients : np.ndarray + gradients : NDArray Gradients corresponding to `trajectory`. - slewrates : np.ndarray + slewrates : NDArray Slew rates corresponding to `trajectory` gradients. """ # Convert normalized trajectory to gradients @@ -411,15 +413,19 @@ def compute_gradients_and_slew_rates( def check_hardware_constraints( - gradients, slewrates, gmax=DEFAULT_GMAX, smax=DEFAULT_SMAX, order=None -): + gradients: NDArray, + slewrates: NDArray, + gmax: float = DEFAULT_GMAX, + smax: float = DEFAULT_SMAX, + order: int | str | None = None, +) -> tuple[bool, float, float]: """Check if a trajectory satisfies the gradient hardware constraints. Parameters ---------- - gradients : np.ndarray + gradients : NDArray Gradients to check - slewrates: np.ndarray + slewrates: NDArray Slewrates to check gmax : float, optional Maximum gradient amplitude in T/m. The default is DEFAULT_GMAX. @@ -450,12 +456,12 @@ def check_hardware_constraints( ########### -def initialize_tilt(tilt, nb_partitions=1): +def initialize_tilt(tilt: str | float | None, nb_partitions: int = 1) -> float: r"""Initialize the tilt angle. Parameters ---------- - tilt : str or float + tilt : str | float | None Tilt angle in rad or name of the tilt. nb_partitions : int, optional Number of partitions. The default is 1. @@ -493,12 +499,12 @@ def initialize_tilt(tilt, nb_partitions=1): raise NotImplementedError(f"Unknown tilt name: {tilt}") -def initialize_algebraic_spiral(spiral): +def initialize_algebraic_spiral(spiral: str | float) -> float: """Initialize the algebraic spiral type. Parameters ---------- - spiral : str or float + spiral : str | float Spiral type or spiral power value. Returns @@ -507,16 +513,16 @@ def initialize_algebraic_spiral(spiral): Spiral power value. """ if isinstance(spiral, Real): - return spiral - return Spirals[spiral] + return float(spiral) + return Spirals[str(spiral)] -def initialize_shape_norm(shape): +def initialize_shape_norm(shape: str | float) -> float: """Initialize the norm for a given shape. Parameters ---------- - shape : str or float + shape : str | float Shape name or p-norm value. Returns @@ -525,5 +531,5 @@ def initialize_shape_norm(shape): Shape p-norm value. """ if isinstance(shape, Real): - return shape - return NormShapes[shape] + return float(shape) + return NormShapes[str(shape)]