Skip to content

Creating coordinate systems¤

This documentation is a collection of functions used to create coordinate systems in cryojax's conventions. The most important functions are make_coordinate_grid and make_frequency_grid.

cryojax.coordinates.make_coordinate_grid(shape: tuple[int, ...], grid_spacing: float | Float[np.ndarray, ''] | Float[Array, ''] = 1.0) -> Float[Array, '*shape ndim'] ¤

Create a real-space cartesian coordinate system on a grid.

Arguments:

  • shape: Shape of the grid, with ndim = len(shape).
  • grid_spacing: The grid spacing (i.e. pixel/voxel size), in units of length.

Returns:

A cartesian coordinate system in real space.

cryojax.coordinates.make_frequency_grid(shape: tuple[int, ...], grid_spacing: float | Float[np.ndarray, ''] | Float[Array, ''] = 1.0, half_space: bool = True) -> Float[Array, '*shape ndim'] ¤

Create a fourier-space cartesian coordinate system on a grid. The zero-frequency component is in the corner.

Arguments:

  • shape: Shape of the grid, with ndim = len(shape).
  • grid_spacing: The grid spacing (i.e. pixel/voxel size), in units of length.
  • half_space: Return a frequency grid on the half space. shape[-1] is the axis on which the negative frequencies are omitted.

Returns:

A cartesian coordinate system in frequency space.

cryojax.coordinates.make_frequency_slice(shape: tuple[int, int], grid_spacing: float | Float[np.ndarray, ''] | Float[Array, ''] = 1.0, half_space: bool = True) -> Float[Array, '1 {shape[0]} {shape[1]} 3'] ¤

Create a fourier-space cartesian coordinate system on a grid, where zero-frequency component is in the center of the grid.

Warning

In the function make_frequency_grid, the convention is that the grid is returned with the zero frequency component is in the corner. In this function, as mentioned above, frequency slices are returned with the zero frequency component in the center. To convert between the two conventions, run

import jax.numpy as jnp
from cryojax.coordinates import make_frequency_slice

frequency_slice_with_zero_in_center = make_frequency_slice((100, 100)) # Shape (1, 100, 100, 3)
frequency_slice_with_zero_in_corner = jnp.fft.ifftshift(frequency_slice_with_zero_in_center, axes=(1, 2))

The reason for the difference is so that this function can be used to directly pass a frequency_slice to the cryojax.simulator.FourierVoxelGridPotential, which requires that the zero is in the center of the grid.

Arguments:

  • shape: Shape of the frequency slice, e.g. shape = (100, 100).
  • grid_spacing: The grid spacing (i.e. voxel size), in units of length.
  • half_space: Return a frequency slice on the half space. shape[-1] is the axis on which the negative frequencies are omitted.

Returns:

The central, \(q_z = 0\) slice of a 3D frequency grid \((q_x, q_y, q_z)\), where zero-frequency component is in the center of the grid.

cryojax.coordinates.make_1d_coordinate_grid(size: int, grid_spacing: float | Float[np.ndarray, ''] | Float[Array, ''] = 1.0) -> Float[Array, '*shape ndim'] ¤

Create a 1D real-space cartesian coordinate array.

Arguments:

  • size: Size of the coordinate array.
  • grid_spacing: The grid spacing (i.e. pixel/voxel size), in units of length.

Returns:

A 1D cartesian coordinate array in real space.

cryojax.coordinates.make_1d_frequency_grid(size: int, grid_spacing: float | Float[np.ndarray, ''] | Float[Array, ''] = 1.0, half_space: bool = True) -> Float[Array, '*shape ndim'] ¤

Create a 1D fourier-space cartesian coordinate array. If half_space = False, the zero-frequency component is in the beginning.

Arguments¤
  • size: Size of the coordinate array.
  • grid_spacing: The grid spacing (i.e. pixel/voxel size), in units of length.
  • half_space: Return a frequency grid on the half space. shape[-1] is the axis on which the negative frequencies are omitted.

Returns:

A 1D cartesian coordinate array in frequency space.