ptyrax.spatial

ptyrax.spatial#

Functions

`R_x`(angle, **kwargs)	Construct a 3x3 rotation matrix for rotation about the x-axis.
`R_y`(angle, **kwargs)	Construct a 3x3 rotation matrix for rotation about the y-axis.
`R_z`(angle, **kwargs)	Construct a 3x3 rotation matrix for rotation about the z-axis.
`angle_safe_jvp`(primals, tangents)
`are_in_plane`(a, b, c)	Check whether three vectors are coplanar (scalar triple product close to zero).
`are_mutually_orthogonal`(*args)	Check whether all given vectors are mutually orthogonal.
`convert_coordinates_into_indices`(...)	Map physical coordinates to array indices on a target grid.
`draw_axis_arrow`(axis_label, xy_start, ...)	Draw a labeled arrow on a matplotlib axes to indicate a coordinate axis.
`interpolate_grid_to_grid`(samples[, ...])	Interpolate complex-valued samples from a source grid onto a target grid.
`is_normalized`(a)	Check whether a vector has unit norm.
`is_orthogonal`(a, b)	Check whether two vectors are orthogonal (dot product close to zero).
`is_parallel`(a, b)	Check whether two vectors are parallel.
`matrix_to_six_dimensional_representation`(...)	Convert a 3x3 rotation matrix to its 6D continuous representation.
`meshgrid`(shape[, pixel_size])	Create a 2D coordinate meshgrid centered at the origin.
`phase_only_exp`(phase)	Compute a unit-magnitude complex exponential $e^{iphi}$ using JAX.
`phase_only_exp_np`(phase)	Compute a unit-magnitude complex exponential $e^{iphi}$ using NumPy.
`plot_geometry`(sample_coordinates, ...[, ...])	Plot the experimental geometry in the x-z plane.
`reflect`(a, n)	Reflect a vector `a` across a plane with unit normal `n`.
`rotate_about_point`(input, rotation, ...)	Rotate a 3D point about a given center of rotation.
`rotation_matrix_from_angles`(angles)	Build a rotation matrix from intrinsic Euler angles (x-y-z order).
`shift_with_interpolation`(x, center, target_shape)	Shift and crop an array using bilinear interpolation.
`shift_with_interpolation_equal_pixel_size`(x, ...)	Extract a shifted sub-region from an array when source and target pixel sizes match.
`shift_with_interpolation_unequal_pixel_size`(x, ...)	Extract a shifted and resampled sub-region when pixel sizes differ.
`six_dimensional_representation_to_matrix`(...)	Converts a 6D rotation representation to a 3x3 rotation matrix.
`three_dimensional_representation_to_matrix`(...)	Converts a 3D rotation representation (with fixed z-axis) to a 3x3 rotation matrix.

Classes

`CoordinateSystem`([rotation, translation, ...])	A coordinate system defined by a translation and rotation in CXI coordinates.
`Rotation`([_representation_6d])	An equinox module wrapping a rotation using the 6D continuous representation.
`SamplingGrid`(shape, pixel_size[, ...])	A discrete sampling grid with pixel sizes and optional FFT-shift awareness.
`YOnlyRotation`(representation_initializer)	A rotation constrained to rotate only about the y-axis.

class ptyrax.spatial.CoordinateSystem(rotation=Rotation(_representation_6d=f32[6]), translation=Array([0., 0., 0.], dtype=float32), model_in_local_frame=False, **kwargs)[source]#

Bases: Module

A coordinate system defined by a translation and rotation in CXI coordinates.

Represents a rigid-body transformation (rotation + translation) of a local coordinate frame relative to the global CXI frame. The CXI convention defines z along the incoming beam direction and y vertically.

The translation can be stored in either the global frame or the local frame, controlled by model_in_local_frame. When model_in_local_frame=True, the internal translation is stored in the local (rotated) frame and converted to global coordinates on access via the translation property.

Variables:

rotation – The Rotation of this coordinate system.
model_in_local_frame – If True, the internal translation is stored in the rotated local frame rather than the global frame.

Parameters:

rotation (Rotation)
translation (Float[Array, '* 3'])
model_in_local_frame (bool)

Example

>>> import jax.numpy as jnp
>>> from ptyrax.spatial import CoordinateSystem, Rotation
>>> cs = CoordinateSystem(
...     rotation=Rotation.from_matrix(jnp.eye(3)),
...     translation=jnp.array([1.0, 0.0, 0.0]),
... )
>>> cs.translation

model_in_local_frame: bool = False#

rotation: Rotation#

property shape: tuple[int, ...]#: Batch shape of the coordinate system.

property translation: Float[Array, '* 3']#

Translation vector in the global CXI frame.

When model_in_local_frame=True, the stored local-frame translation is rotated into the global frame before being returned.

property translation_internal: Float[Array, '* 3']#

Translation vector in the local (rotated) frame.

When model_in_local_frame=False, the stored global-frame translation is rotated into the local frame before being returned.

property x_axis: Float[Array, '* 3']#: Unit vector along the local x-axis, expressed in global coordinates.

property y_axis: Float[Array, '* 3']#: Unit vector along the local y-axis, expressed in global coordinates.

property z_axis: Float[Array, '* 3']#: Unit vector along the local z-axis, expressed in global coordinates.

ptyrax.spatial.R_x(angle, **kwargs)[source]#

Construct a 3x3 rotation matrix for rotation about the x-axis.

Parameters:: angle (float) – Rotation angle in degrees.
Returns:: A (3, 3) rotation matrix.
Return type:: Float[Array, ‘3 3’]

ptyrax.spatial.R_y(angle, **kwargs)[source]#

Construct a 3x3 rotation matrix for rotation about the y-axis.

Parameters:: angle (float) – Rotation angle in degrees.
Returns:: A (3, 3) rotation matrix.
Return type:: Float[Array, ‘3 3’]

ptyrax.spatial.R_z(angle, **kwargs)[source]#

Construct a 3x3 rotation matrix for rotation about the z-axis.

Parameters:: angle (float) – Rotation angle in degrees.
Returns:: A (3, 3) rotation matrix.
Return type:: Float[Array, ‘3 3’]

class ptyrax.spatial.Rotation(_representation_6d=<factory>)[source]#

Bases: Module

An equinox module wrapping a rotation using the 6D continuous representation.

Stores rotations internally as a 6D vector (the first two rows of the rotation matrix), which provides a continuous, singularity-free parameterization suitable for gradient-based optimization. Based on Zhou et al. (NeurIPS 2019).

Supports batched rotations via leading batch dimensions, composition via *, inversion via the inv property, and indexing via [].

Example

>>> import jax.numpy as jnp
>>> from ptyrax.spatial import Rotation
>>> rot = Rotation.from_matrix(jnp.eye(3))
>>> rot.as_matrix()  # Returns identity matrix

Parameters:: _representation_6d (Float[Array, '* 6'])

as_matrix()[source]#

Convert this rotation to a 3x3 rotation matrix.

Returns:: A shape (*, 3, 3) orthonormal rotation matrix.
Return type:: Float[Array, ‘* 3 3’]

as_scipy_rotation()[source]#

Convert this rotation to a JAX scipy Rotation object.

Returns:: A jax.scipy.spatial.transform.Rotation instance.
Return type:: Rotation

classmethod from_matrix(matrix)[source]#

Construct a rotation from a 3x3 rotation matrix.

Parameters:: matrix (Float[Array, '* 3 3']) – A shape (*, 3, 3) rotation matrix.
Returns:: A new Rotation instance.
Return type:: Self

property inv: Self#: The inverse rotation (transpose of the rotation matrix).

property shape: tuple[int, ...]#: Batch shape of the rotation.

class ptyrax.spatial.SamplingGrid(shape, pixel_size, origin_shift=None, fftshifted=False)[source]#

Bases: Module

A discrete sampling grid with pixel sizes and optional FFT-shift awareness.

Represents a regular 2D or 3D grid centered at the origin (or shifted by origin_shift), with coordinates following CXI convention: x along dimension 0, y along dimension 1, using numpy 'ij' indexing.

Grid coordinates range from -N//2 to (N+1)//2 - 1 along each axis, scaled by the corresponding pixel size. When fftshifted=True, the coordinate arrays are reordered to match FFT output layout.

Variables:

shape – Grid dimensions as (n_x, n_y) or (n_x, n_y, n_z).
pixel_size – Physical size of each pixel along each dimension, shape (*, d).
origin_shift – Optional shift of the grid origin, shape (d,).
fftshifted – If True, coordinate arrays are fftshifted.

Parameters:

shape (tuple[int, ...])
pixel_size (Float[Array, '* d'] | ndarray)
origin_shift (Float[Array, 'd'] | ndarray)
fftshifted (bool)

Example

>>> import jax.numpy as jnp
>>> from ptyrax.spatial import SamplingGrid
>>> grid = SamplingGrid(shape=(64, 64), pixel_size=jnp.array([1e-6, 1e-6]))
>>> grid.x  # 1D coordinates along x
>>> grid.meshgrid  # Full 2D coordinate meshgrid

__plot__(fig=None, gs=None, show=True, **kwargs)[source]#

Parameters:

fig (Figure)
gs (SubplotSpec)
show (bool)

Return type:

tuple[Figure, SubplotSpec, list[Line2D]]

property bounds: Float[Array, '... d 2']#

Coordinate bounds as a (*, d, 2) array of [min, max] per axis.

If origin_shift is set, it is added to the bounds.

edges(n_per_edge=100)[source]#

Generate 3D coordinates along the four edges of the 2D grid boundary.

Returns points tracing the rectangular boundary of the grid in the local x-y plane (z=0), useful for visualization.

Parameters:: n_per_edge (int) – Number of sample points along each edge.
Returns:: An (N, 3) array of boundary coordinates with z=0.
Return type:: Float[Array, ‘N 3’]

property extent: Float[Array, '... 2d']#: Flattened bounds as a (*, 2d) array, e.g. (x_min, x_max, y_min, y_max).

fftshifted: bool = False#

property field_of_view: Float[Array, 'd 2']#: Alias for bounds.

classmethod from_tuples(shape, pixel_size, origin_shift=None, fftshifted=False)[source]#

Construct a SamplingGrid from plain tuples.

Convenience constructor that converts tuple arguments to the appropriate array types.

Parameters:

shape (tuple[int | float]) – Grid dimensions, e.g. (64, 64).
pixel_size (tuple[float]) – Physical pixel size per dimension, e.g. (1e-6, 1e-6).
origin_shift (tuple[float]) – Optional origin offset per dimension.
fftshifted (bool) – Whether the grid coordinates should be FFT-shifted.

Returns:

A new SamplingGrid instance.

Return type:

SamplingGrid

lines_for_plot(n_x=20, n_y=20)[source]#

Generate grid lines for 2D plotting.

Produces horizontal and vertical line segments spanning the grid, suitable for overlaying on a plot.

Parameters:

n_x (int) – Number of vertical grid lines.
n_y (int) – Number of horizontal grid lines.

Returns:

An array of line segments, each of shape (2, N).

Return type:

Float[Array, ‘N 2’]

property meshgrid: Float[Array, 'd N_x N_y'] | Float[Array, 'd N_x N_y N_z']#: Full coordinate meshgrid with shape (d, N_x, N_y) or (d, N_x, N_y, N_z).

property n_x: int#: Number of grid points along the x-axis.

property n_x_max: int#: Maximum x index (inclusive) in the centered grid.

property n_x_min: int#: Minimum x index in the centered grid.

property n_y: int#: Number of grid points along the y-axis.

property n_y_max: int#: Maximum y index (inclusive) in the centered grid.

property n_y_min: int#: Minimum y index in the centered grid.

property n_z: int#

Number of grid points along the z-axis.

Raises:: ValueError – If the grid is 2D.

property n_z_max: int#: Maximum z index (inclusive) in the centered grid.

property n_z_min: int#: Minimum z index in the centered grid.

property ndim: int#: Number of spatial dimensions (2 or 3).

origin_shift: Float[Array, 'd'] | ndarray = None#

pixel_size: Float[Array, '* d'] | ndarray#

property rr: Float[Array, 'N_x N_y'] | Float[Array, 'N_x N_y N_z']#: Radial distance from the origin at each grid point.

property rr2: Float[Array, 'N_x N_y'] | Float[Array, 'N_x N_y N_z']#: Squared radial distance from the origin at each grid point.

shape: tuple[int, ...]#

subgrid_coordinates(n_x=20, n_y=20)[source]#

Generate a coarse subgrid of 3D coordinates spanning the grid extent.

Useful for visualization or sampling a sparse set of points across the grid.

Parameters:

n_x (int) – Number of points along the x-axis.
n_y (int) – Number of points along the y-axis.

Returns:

An (n_x * n_y, 3) array of coordinates with z=0.

Return type:

Float[Array, ‘n_x n_y 3’]

to_far_field(wavelength=1.0, propagation_distance=1.0, fftshifted=False)[source]#

Compute the reciprocal-space (far-field) grid corresponding to this grid.

Applies the Fraunhofer diffraction relation to convert pixel sizes from real space to reciprocal space:

$$Delta q_i = frac{lambda cdot z}{Delta x_i cdot N_i}$$

where $lambda$ is the wavelength, $z$ is the propagation distance, $Delta x_i$ is the real-space pixel size, and $N_i$ is the number of pixels along axis $i$.

Parameters:

wavelength (Float[Array, ''] | float) – Radiation wavelength (scalar).
propagation_distance (Float[Array, ''] | float) – Propagation distance to the far-field plane (scalar).
fftshifted (bool) – Whether the returned grid should be FFT-shifted.

Returns:

A new SamplingGrid with reciprocal-space pixel sizes.

Raises:

ValueError – If wavelength is not scalar.

Return type:

Self

property x: Float[Array, 'n_x']#: 1D array of physical x-coordinates.

property x_bounds: Float[Array, '... 2']#: Min and max physical x-coordinates as a (2,) array.

property x_max: float#: Maximum physical x-coordinate.

property x_min: float#: Minimum physical x-coordinate.

property x_pixel_size: float#: Physical pixel size along the x-axis.

property xx: Float[Array, 'N_x N_y'] | Float[Array, 'N_x N_y N_z']#: X-component of the coordinate meshgrid.

property y: Float[Array, 'n_y']#: 1D array of physical y-coordinates.

property y_bounds: Float[Array, '... 2']#: Min and max physical y-coordinates as a (2,) array.

property y_max: float#: Maximum physical y-coordinate.

property y_min: float#: Minimum physical y-coordinate.

property y_pixel_size: float#: Physical pixel size along the y-axis.

property yy: Float[Array, 'N_x N_y'] | Float[Array, 'N_x N_y N_z']#: Y-component of the coordinate meshgrid.

property z: Float[Array, 'n_z']#: 1D array of physical z-coordinates.

property z_bounds: Float[Array, '... 2']#: Min and max physical z-coordinates as a (2,) array.

property z_max: float#: Maximum physical z-coordinate.

property z_min: float#: Minimum physical z-coordinate.

property z_pixel_size: float#

Physical pixel size along the z-axis.

Raises:: ValueError – If the grid is 2D.

property zz: Float[Array, 'N_x N_y N_z'] | Float[Array, 'N_x N_y N_z']#: Z-component of the coordinate meshgrid (3D grids only).

class ptyrax.spatial.YOnlyRotation(representation_initializer)[source]#

Bases: Rotation

A rotation constrained to rotate only about the y-axis.

Useful for samples in simple reflection geometries where only a rotation around the vertical (y) axis are presumed to be present. The y-axis component of the 6D representation is fixed to [0, 1, 0].

Parameters:: representation_initializer (Callable[[], Float[Array, '3']]) – A callable returning a shape (3,) array representing the x-axis direction of the rotation.
Return type:: Self

as_matrix()#

Convert this rotation to a 3x3 rotation matrix.

Returns:: A shape (*, 3, 3) orthonormal rotation matrix.
Return type:: Float[Array, ‘* 3 3’]

as_scipy_rotation()#

Convert this rotation to a JAX scipy Rotation object.

Returns:: A jax.scipy.spatial.transform.Rotation instance.
Return type:: Rotation

classmethod from_matrix(matrix)#

Construct a rotation from a 3x3 rotation matrix.

Parameters:: matrix (Float[Array, '* 3 3']) – A shape (*, 3, 3) rotation matrix.
Returns:: A new Rotation instance.
Return type:: Self

property inv: Self#: The inverse rotation (transpose of the rotation matrix).

representation#: alias of Float[Array, '3']

property representation_6d: Float[Array, '6']#

property shape: tuple[int, ...]#: Batch shape of the rotation.

ptyrax.spatial.angle_safe_jvp(primals, tangents)[source]#

Parameters:

primals (Complex[Array, '...'])
tangents (Complex[Array, '...'])

Return type:

tuple[Float[Array, ‘…’], Float[Array, ‘…’]]

ptyrax.spatial.are_in_plane(a, b, c)[source]#

Check whether three vectors are coplanar (scalar triple product close to zero).

Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)
b (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)
c (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)

Return type:

bool

ptyrax.spatial.are_mutually_orthogonal(*args)[source]#

Check whether all given vectors are mutually orthogonal.

Parameters:: *args (tuple[Array | ndarray | bool | number | bool | int | float | complex | LiteralArray, ...]) – Two or more vectors to test.
Returns:: True if every pair of vectors is orthogonal.
Return type:: bool

ptyrax.spatial.convert_coordinates_into_indices(coordinates, target_sampling)[source]#

Map physical coordinates to array indices on a target grid.

Converts physical (x, y) coordinates into floating-point array indices relative to the target SamplingGrid. Also returns a boolean mask indicating which coordinates fall within the grid boundaries.

Parameters:

coordinates (Float[Array, '... 2']) – Physical coordinates with shape (..., 2) where the leading dimension indexes (x, y).
target_sampling (SamplingGrid) – The target grid defining pixel sizes and shape.

Returns:

indices: Floating-point array indices, shape (2, ...).
mask: Boolean array indicating in-bounds coordinates.

Return type:

A tuple of

ptyrax.spatial.draw_axis_arrow(axis_label, xy_start, xy_end, ax, **kwargs)[source]#

Draw a labeled arrow on a matplotlib axes to indicate a coordinate axis.

Parameters:

axis_label (str) – Text label for the axis (e.g. '$x_s$').
xy_start (tuple[float, float]) – (x, y) coordinates of the arrow tip.
xy_end (tuple[float, float]) – (x, y) coordinates of the arrow tail.
ax (Axes) – The matplotlib axes to draw on.
**kwargs – Additional keyword arguments passed to ax.annotate.

Returns:

The axes with the arrow drawn.

Return type:

Axes

ptyrax.spatial.interpolate_grid_to_grid(samples, source_grid=None, source_coordinate_system=CoordinateSystem(rotation=Rotation(_representation_6d=f32[6]), _translation=f32[3]), target_grid=None, target_coordinate_system=CoordinateSystem(rotation=Rotation(_representation_6d=f32[6]), _translation=f32[3]), interpolation_mode=None, equal_pixel_size=False)[source]#

Interpolate complex-valued samples from a source grid onto a target grid.

Handles both equal and unequal pixel sizes between source and target grids, accounting for the relative translation between their coordinate systems. The translation is projected into the source’s local frame to determine the pixel-space offset.

Parameters:

samples (Complex[Array, '* m0 n0']) – Complex array of shape (*, m0, n0) on the source grid.
source_grid (SamplingGrid) – The SamplingGrid of the source data.
source_coordinate_system (CoordinateSystem) – The CoordinateSystem of the source grid.
target_grid (SamplingGrid) – The SamplingGrid to interpolate onto. Defaults to source_grid if not provided.
target_coordinate_system (CoordinateSystem) – The CoordinateSystem of the target grid.
interpolation_mode (Literal['real_imaginary', 'amplitude_phase', 'sliced_rounded', 'sliced_interpolated']) – Interpolation strategy. See shift_with_interpolation_equal_pixel_size() and shift_with_interpolation_unequal_pixel_size().
equal_pixel_size (bool) – If True, use the faster equal-pixel-size path.

Returns:

The interpolated samples on the target grid, shape (*, m1, n1).

Raises:

ValueError – If source_grid is not provided.

Return type:

Complex[Array, ‘* m1 n1’]

ptyrax.spatial.is_normalized(a)[source]#

Check whether a vector has unit norm.

Parameters:: a (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)
Return type:: bool

ptyrax.spatial.is_orthogonal(a, b)[source]#

Check whether two vectors are orthogonal (dot product close to zero).

Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)
b (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)

Return type:

bool

ptyrax.spatial.is_parallel(a, b)[source]#

Check whether two vectors are parallel.

Parameters:

a (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)
b (Array | ndarray | bool | number | bool | int | float | complex | LiteralArray)

Return type:

bool

ptyrax.spatial.matrix_to_six_dimensional_representation(rotation_matrix)[source]#

Convert a 3x3 rotation matrix to its 6D continuous representation.

Uses the first two rows of the rotation matrix as the 6D representation, following the approach from Zhou et al. (NeurIPS 2019) for continuous rotation representations suitable for gradient-based optimization.

Parameters:: rotation_matrix (Float[Array, '... 3 3']) – A shape (..., 3, 3) rotation matrix.
Returns:: A shape (..., 6) array containing the first two rows of the rotation matrix concatenated.
Return type:: Float[Array, ‘… 6’]

ptyrax.spatial

Contents

ptyrax.spatial#