Volume projection and integration onto the plane

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

integrate(volume_representation: ~VolRep, 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}']

cryojax.simulator.AutoVolumeProjection(cryojax.simulator.AbstractVolumeIntegrator)

Volume projection auto selection from cryoJAX AbstractVolumeIntegrator implementations.

Info

Based on the cryojax.simulator.AbstractVolumeRepresentation passed at runtime, this class chooses a default projection method. In particular,

Volume representation	Projection method	Atom or voxel?
cryojax.simulator.GaussianMixtureVolume	cryojax.simulator.GaussianMixtureProjection	atom
cryojax.simulator.IndependentAtomVolume	cryojax.simulator.FFTAtomProjection	atom
cryojax.simulator.FourierVoxelGridVolume or cryojax.simulator.FourierVoxelSplineVolume	cryojax.simulator.FourierSliceExtraction	voxel
cryojax.simulator.RealVoxelGridVolume	cryojax.simulator.RealVoxelProjection	voxel

To use advanced options for a given projection method, see each respective class.

Warning

If using cryojax.simulator.FFTAtomRenderFn or cryojax.simulator.RealVoxelProjection, jax-finufft must be installed. See the cryoJAX installation instructions for installing jax-finufft.

__init__

__init__()

Arguments:

• options: Keyword arguments passed to the resolved projection method, e.g.GaussianMixtureProjection(**options)`.

integrate

integrate(volume_representation: cryojax.simulator.AbstractVolumeRepresentation, 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}']

Automatically select volume projection method given a volume representation.

Arguments:

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

Returns:

The volume projection in real or Fourier space at the AbstractImageConfig.padded_shape and the image_config.pixel_size.

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__

__init__(*, outputs_integral: bool = True, 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.
• 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

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 image configuration.
• outputs_real_space: If True, return the image in real space. Otherwise, return in Fourier.

Returns:

The volume projection in real or Fourier space at the AbstractImageConfig.padded_shape and the image_config.pixel_size.

---

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

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

__init__

__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

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 image configuration.
• outputs_real_space: If True, return the image in real space. Otherwise, return in Fourier.

Returns:

The volume projection in real or Fourier space, at the image_config.padded_shape.

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, sampling_mode: Literal['average', 'point'] = 'average', use_error_functions: bool = False, n_batches: int = 1)

__init__

__init__(*, upsampling_factor: int | None = None, shape: tuple[int, int] | None = None, sampling_mode: Literal['average', 'point'] = 'average', use_error_functions: bool = False, 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.
• sampling_mode: If 'average', use error functions to sample the projected volume at a pixel to be the average value using gaussian integrals. If 'point', the volume at a pixel will be evaluated by evaluating the gaussian at a point.
• 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

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 image configuration.
• outputs_real_space: If True, return the image in real space. Otherwise, return in Fourier.

Returns:

The volume projection in real or Fourier space at the AbstractImageConfig.padded_shape and the image_config.pixel_size.

---

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

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

__init__

__init__(*, sampling_mode: Literal['average', 'point'] = 'average', upsample_factor: int | None = None, shape: tuple[int, int] | None = None, eps: float = 1e-06, opts: Any = None)

Arguments:

• sampling_mode: If 'average', convolve with a box function to sample the projected volume at a pixel to be the average value of the underlying continuous function. If 'point', the volume at a pixel will be point sampled.
• 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

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 volume projection in real or Fourier space at the AbstractImageConfig.padded_shape and the image_config.pixel_size.