Skip to content

Volume integration methods¤

cryojax provides different methods for integrating volumes onto a plane to generate an image.

cryojax.simulator.AbstractVolumeIntegrator

cryojax.simulator.AbstractVolumeIntegrator ¤

Base class for a method of integrating a volume onto the exit plane.

integrate(volume_representation: ~VolRepT, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}'] | Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] ¤

Integration methods for voxel-based structures¤

cryojax.simulator.FourierSliceExtraction(cryojax.simulator.AbstractVolumeIntegrator) ¤

Integrate points to the exit plane using the Fourier projection-slice theorem.

This extracts slices using interpolation methods housed in cryojax.ndimage.map_coordinates and cryojax.ndimage.map_coordinates_spline.

__init__(*, outputs_integral: bool = True, correction_mask: cryojax.ndimage.transforms.InverseSincMask | None = None, out_of_bounds_mode: str = 'fill', fill_value: complex = 0j) ¤

Arguments:

  • outputs_integral: If True, return the fourier slice multiplied by the voxel size. Including the voxel size numerical approximates the projection integral and is necessary for simulating images in physical units.
  • correction_mask: A cryojax.ndimage.transforms.InverseSincMask for performing sinc-correction on the linear-interpolated projections. This should be computed on a coordinate grid with shape matching the FourierVoxelGridVolume.shape.
  • out_of_bounds_mode: Specify how to handle out of bounds indexing. See cryojax.ndimage.map_coordinates for documentation.
  • fill_value: Value for filling out-of-bounds indices. Used only when out_of_bounds_mode = "fill".
integrate(volume_representation: cryojax.simulator.FourierVoxelGridVolume | cryojax.simulator.FourierVoxelSplineVolume, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] ¤

Integrate the volume at the AbstractImageConfig settings of a voxel-based representation in fourier-space, using fourier slice extraction.

Arguments:

  • volume_representation: The volume representation.
  • image_config: The configuration of the resulting image.
  • outputs_real_space: If True, return the image in real space. Otherwise, return in fourier.

Returns:

The extracted fourier voxels of the volume_representation, at the image_config.padded_shape and the image_config.pixel_size.

extract_fourier_slice_from_spline(spline_coefficients: Complex[Array, 'coeff_dim coeff_dim coeff_dim'], frequency_slice_in_pixels: Float[Array, '1 dim dim 3']) -> Complex[Array, 'dim dim//2+1'] ¤

Extract a fourier slice using the interpolation defined by spline_coefficients at coordinates frequency_slice_in_pixels.

Arguments:

  • spline_coefficients: Spline coefficients of the density grid in fourier space. The coefficients should be computed from a fourier_voxel_grid with the zero frequency component in the center. These are typically computed with the function cryojax.image.compute_spline_coefficients.
  • frequency_slice_in_pixels: Frequency central slice coordinate system. The zero frequency component should be in the center.
  • voxel_size: The voxel size of the fourier_voxel_grid. This argument is not used in the FourierSliceExtraction class.
  • wavelength_in_angstroms: The wavelength of the incident electron beam. This argument is not used in the FourierSliceExtraction class.

Returns:

The interpolated fourier slice at coordinates frequency_slice_in_pixels.

extract_fourier_slice_from_grid(fourier_voxel_grid: Complex[Array, 'dim dim dim'], frequency_slice_in_pixels: Float[Array, '1 dim dim 3']) -> Complex[Array, 'dim dim//2+1'] ¤

Extract a fourier slice of the fourier_voxel_grid at coordinates frequency_slice_in_pixels.

Arguments:

  • fourier_voxel_grid: Density grid in fourier space. The zero frequency component should be in the center.
  • frequency_slice_in_pixels: Frequency central slice coordinate system. The zero frequency component should be in the center.
  • voxel_size: The voxel size of the fourier_voxel_grid. This argument is not used in the FourierSliceExtraction class.
  • wavelength_in_angstroms: The wavelength of the incident electron beam. This argument is not used in the FourierSliceExtraction class.

Returns:

The interpolated fourier slice at coordinates frequency_slice_in_pixels.

cryojax.simulator.RealVoxelProjection(cryojax.simulator.AbstractVolumeIntegrator) ¤

Integrate points onto the exit plane using non-uniform FFTs.

__init__(*, eps: float = 1e-06, opts: Any = None) ¤

Arguments:

  • eps: See jax-finufft for documentation.
  • opts: A jax_finufft.options.Opts or jax_finufft.options.NestedOpts dataclass. See jax-finufft for documentation.
integrate(volume_representation: cryojax.simulator.RealVoxelGridVolume, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] ¤

Integrate the volume at the AbstractImageConfig settings of a voxel-based representation in real-space, using non-uniform FFTs.

Arguments:

  • volume_representation: The volume representation.
  • image_config: The configuration of the resulting image.
  • outputs_real_space: If True, return the image in real space. Otherwise, return in fourier.

Returns:

The projection integral of the volume_representation in fourier space, at the image_config.padded_shape and the image_config.pixel_size.

Integration methods for atom-based based structures¤

cryojax.simulator.GaussianMixtureProjection(cryojax.simulator.AbstractVolumeIntegrator) ¤

GaussianMixtureProjection(*, upsampling_factor: int | None = None, shape: tuple[int, int] | None = None, use_error_functions: bool = True, n_batches: int = 1)

__init__(*, upsampling_factor: int | None = None, shape: tuple[int, int] | None = None, use_error_functions: bool = True, n_batches: int = 1) ¤

Arguments:

  • shape: The shape of the plane on which projections are computed before padding or cropping to the AbstractImageConfig.padded_shape. This argument is particularly useful if the AbstractImageConfig.padded_shape is much larger than the protein.
  • use_error_functions: If True, use error functions to evaluate the projected volume at a pixel to be the average value within the pixel using gaussian integrals (i.e. apply antialiasing). If False, the volume at a pixel will be evaluated as a gaussian.
  • n_batches: The number of batches over groups of positions used to evaluate the projection. By default, n_batches = 1, which computes a projection for all positions at once. This is useful to decrease GPU memory usage.
integrate(volume_representation: cryojax.simulator.GaussianMixtureVolume, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] ¤

Compute a projection from gaussians.

Arguments:

  • volume_representation: The volume representation.
  • image_config: The configuration of the resulting image.
  • outputs_real_space: If True, return the image in real space. Otherwise, return in fourier.

Returns:

The integrated volume in real or fourier space at the AbstractImageConfig.padded_shape.

cryojax.simulator.FFTAtomProjection(cryojax.simulator.AbstractVolumeIntegrator) ¤

Integrate atomic parametrization of a volume onto the exit plane using non-uniform FFTs plus convolution.

__init__(*, antialias: bool = True, upsample_factor: int | None = None, shape: tuple[int, int] | None = None, eps: float = 1e-06, opts: Any = None) ¤

Arguments:

  • antialias: If True, apply an anti-aliasing filter to more accurately sample the volume.
  • upsample_factor: If provided, first compute an upsampled version of the image at pixel size image_config.pixel_size / upsample_factor. Then, downsample with cryojax.ndimage.block_reduce_downsample to locally average to the correct pixel size. This is useful for reducing aliasing.
  • shape: If given, first compute the image at shape, then pad or crop to image_config.padded_shape.
  • eps: See jax-finufft for documentation.
  • opts: A jax_finufft.options.Opts or jax_finufft.options.NestedOpts dataclass. See jax-finufft for documentation.
integrate(volume_representation: cryojax.simulator.IndependentAtomVolume, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim//2+1}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] ¤

Compute a projection from scattering factors per atom type from the IndependentAtomVolume.

Arguments:

  • volume_representation: The volume representation.
  • image_config: The configuration of the resulting image.
  • outputs_real_space: If True, return the image in real space. Otherwise, return in fourier.

Returns:

The integrated volume in real or fourier space at the AbstractImageConfig.padded_shape.