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(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.GaussianFourierVolume	cryojax.simulator.GaussianFourierProjection	atom
cryojax.simulator.FourierVoxelGridVolume	cryojax.simulator.FourierSliceExtraction	voxel

Note that cryojax.simulator.RealVoxelGridVolume does not have an associated projection method.

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

__init__()

Arguments:

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

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.

The interpolation method is read from FourierVoxelGridVolume.interp.

__init__(*, boundary: str = 'fill', unroll: bool | Literal['auto'] = 'auto')

Arguments:

• boundary: What to return for frequencies outside the fourier box. See cryojax.ndimage.sample_fft_slice.
• unroll: Passed to cryojax.ndimage.sample_fft_slice. With "auto" (the default), this is True for interp="cubic" and False otherwise.

integrate(volume_representation: cryojax.simulator.FourierVoxelGridVolume, 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.EwaldSphereExtraction(cryojax.simulator.AbstractVolumeIntegrator)

Integrate points to the exit plane by extracting a surface of the ewald sphere in fourier space.

The interpolation method is read from FourierVoxelGridVolume.interp.

__init__(*, boundary: str = 'fill', unroll: bool | Literal['auto'] = 'auto')

Arguments:

• boundary: What to return for frequencies outside the fourier box. See cryojax.ndimage.sample_fft_slice.
• unroll: Passed to cryojax.ndimage.sample_fft_slice. With "auto" (the default), this is True for interp="cubic" and False otherwise.

integrate(volume_representation: cryojax.simulator.FourierVoxelGridVolume, image_config: cryojax.simulator.AbstractImageConfig, outputs_real_space: bool = False) -> Complex[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}'] | Float[Array, '{image_config.padded_y_dim} {image_config.padded_x_dim}']

Extract the ewald sphere surface.

Arguments:

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

Returns:

The Ewald sphere surface in the real-space or fourier-space at the image_config.padded_shape, image_config.pixel_size, and image_config.voltage_in_kilovolts.

Integration methods for atom-based based structures

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

Speed up gradients with Pallas

integrator = cxs.GaussianMixtureProjection(
    n_spread=7, enable_pallas={"bwd": True}
)

__init__(*, shape: tuple[int, int] | None = None, sampling_mode: Literal['average', 'point'] = 'average', n_batches: int = 1, n_spread: int | tuple[int, ...] | None = None, enable_pallas: bool | Mapping[str, bool] | None = None)

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. Applies to both the dense (n_spread=None) and spreading (n_spread set) backends.
• n_spread: If None (default), compute the projection with dense gaussian integrals evaluated over the whole grid. If an int, instead directly spread each gaussian onto only the n_spread nearest grid points (per dimension), trading accuracy for speed. If a tuple of ints (one value per gaussian component, i.e. of length GaussianMixtureVolume.amplitudes.shape[-1]), spread each gaussian component with its own width instead of one shared width -- useful when a volume's gaussian components have widths spanning an order of magnitude or more (e.g. X-ray/electron scattering factors written as a sum of 5 gaussians), where a single n_spread would either truncate the widest components or waste computation spreading the narrowest ones too widely. See cryojax.simulator.suggest_n_spread to choose these values from volume_representation.variances.
• enable_pallas: Use the Pallas/Triton GPU backend instead of pure-JAX for the n_spread spreading backend (ignored if n_spread is None). Pallas is JAX's framework for writing custom GPU/TPU kernels. This is most advantageous for the backward pass -- {"bwd": True}. See cryojax.ndimage.spread_gaussians_2d's enable_pallas for the full picture. None (default) defers to CRYOJAX_ENABLE_PALLAS.

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.GaussianFourierProjection(cryojax.simulator.AbstractVolumeIntegrator)

Integrate atomic parametrization of a volume onto the exit plane from a GaussianFourierVolume using non-uniform FFTs and Fourier-domain convolution. Good when kernels span at least a couple pixels; see cryojax.simulator.GaussianMixtureProjection for an alternative that directly spreads narrow kernels onto the grid.

Info

By default, the non-uniform FFT runs on a pure-JAX backend using nufftax. Setting the environment variable CRYOJAX_FINUFFT_BACKEND=jax-finufft switches to jax-finufft, which can be more computationally efficient and less memory-demanding, at the cost of being trickier to install and having more limited integration with multi-GPU JAX.

__init__(*, sampling_mode: Literal['average', 'point'] = 'average', upsample_factor: int | float = 1.0, eps: float = 1e-06, shape: tuple[int, int] | None = None, options: dict[str, Any] = {})

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: How much to upsample the grid on which atoms are spread onto.
• eps: The precision of the underlying non-uniform FFT implementation. See finufft for documentation.
• shape: If given, first compute the image at shape, then pad or crop to image_config.padded_shape.
• options: A dictionary of options for advanced usage, passed directly to the underlying non-uniform FFT implementation.

integrate(volume_representation: cryojax.simulator.GaussianFourierVolume, 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 GaussianFourierVolume.

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.