Skip to content

Modeling cryo-EM volumes¤

There are many different volume representations of biological structures for cryo-EM, including atomic models, voxel maps, and neural network representations. Further, there are many ways to generate these volumes, such as from protein generative modeling and molecular dynamics. The optimal implementation to use depends on the user's needs. Therefore, CryoJAX supports a variety of these representations as well as a modeling interface for usage downstream. This page discusses how to use this interface and documents the volumes included in the library.

Core base classes¤

cryojax.simulator.AbstractVolumeParametrisation

cryojax.simulator.AbstractVolumeParametrisation ¤

Abstract interface for a parametrisation of a volume.

cryojax.simulator.AbstractVolumeParametrisation.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation abstractmethod ¤

Core interface for computing the representation of the volume.

cryojax.simulator.AbstractVolumeRepresentation

cryojax.simulator.AbstractVolumeRepresentation ¤

Abstract interface for the representation of a volume, such as atomic coordinates, voxels, or a neural network.

cryojax.simulator.AbstractVolumeRepresentation.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self abstractmethod ¤

Rotate the coordinate system of the volume.

Volume representations¤

Point clouds¤

cryojax.simulator.AbstractPointCloudVolume

cryojax.simulator.AbstractPointCloudVolume ¤

Abstract interface for a volume represented as a point-cloud.

cryojax.simulator.AbstractPointCloudVolume.translate_to_pose(pose: AbstractPose) -> Self abstractmethod ¤

cryojax.simulator.GaussianMixtureVolume ¤

A representation of a volume as a mixture of gaussians, with multiple gaussians used per position.

The convention of allowing multiple gaussians per position follows "Robust Parameterization of Elastic and Absorptive Electron Atomic Scattering Factors" by Peng et al. (1996), where amplitudes follows \(a_i\) and \(b_i\) follows variances (multiplied by \(8\pi^2\) to convert to a variance).

cryojax.simulator.GaussianMixtureVolume.__init__(positions: Float[NDArrayLike, 'n_gaussians 3'], amplitudes: Float[NDArrayLike, 'n_positions n_gaussians'], variances: Float[NDArrayLike, 'n_positions n_gaussians']) ¤

Arguments:

  • positions: The coordinates of the gaussians in units of angstroms.
  • amplitudes: The amplitude for each gaussian. To simulate in physical units of a scattering potential, this should have units of angstroms.
  • variances: The variance for each gaussian. This has units of angstroms squared.
cryojax.simulator.GaussianMixtureVolume.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> Self ¤

Since this class is itself an implementation of an AbstractVolumeParametrisation, this function maps to the identity.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.
cryojax.simulator.GaussianMixtureVolume.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self ¤

Return a new potential with rotated positions.

cryojax.simulator.GaussianMixtureVolume.translate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with rotated positions.

cryojax.simulator.GaussianMixtureVolume.to_real_voxel_grid(shape: tuple[int, int, int], voxel_size: Float[NDArrayLike, ''] | float, *, batch_options: dict[str, Any] = {}) -> Float[Array, '{shape[0]} {shape[1]} {shape[2]}'] ¤

Return a voxel grid of the potential in real space.

Arguments:

  • shape: The shape of the resulting voxel grid.
  • voxel_size: The voxel size of the resulting voxel grid.
  • batch_options: Advanced options for rendering. This is a dictionary with the following keys:
    • "batch_size": The number of z-planes to evaluate in parallel with jax.vmap. By default, 1.
    • "n_batches": The number of iterations used to evaluate the volume, where the iteration is taken over groups of gaussians. This is useful if batch_size = 1 and GPU memory is exhausted. By default, 1.

Returns:

A voxel grid of shape and voxel size voxel_size.

cryojax.simulator.PengIndependentAtomPotential ¤

The scattering potential parameterized as a mixture of five gaussians per atom (Peng et al. 1996).

Info

Use the following to load a PengIndependentAtomPotential from tabulated electron scattering factors

from cryojax.io import read_atoms_from_pdb
from cryojax.simulator import (
    PengIndependentAtomPotential, PengScatteringFactorParameters
)

# Load positions of atoms and one-hot encoded atom names
atom_positions, atom_types = read_atoms_from_pdb(...)
parameters = PengScatteringFactorParameters(atom_types)
potential = PengIndependentAtomPotential.from_tabulated_parameters(
    atom_positions, parameters
)

References:

  • Peng, L-M. "Electron atomic scattering factors and scattering potentials of crystals." Micron 30.6 (1999): 625-648.
  • Peng, L-M., et al. "Robust parameterization of elastic and absorptive electron atomic scattering factors." Acta Crystallographica Section A: Foundations of Crystallography 52.2 (1996): 257-276.
cryojax.simulator.PengIndependentAtomPotential.__init__(atom_positions: Float[NDArrayLike, 'n_atoms 3'], amplitudes: Float[Array, 'n_atoms n_gaussians'], b_factors: Float[Array, 'n_atoms n_gaussians']) ¤

Arguments:

  • atom_positions: The coordinates of the atoms in units of angstroms.
  • amplitudes: The strength for each atom and gaussian per atom. This are \(a_i\) from Peng et al. (1996) and has units of angstroms.
  • b_factors: The B-factors for each atom and gaussian per atom. This are \(b_i\) from Peng et al. (1996) and has units of angstroms squared.
cryojax.simulator.PengIndependentAtomPotential.from_tabulated_parameters(atom_positions: Float[NDArrayLike, 'n_atoms 3'], parameters: PengScatteringFactorParameters, extra_b_factors: Optional[Float[NDArrayLike, ' n_atoms']] = None) -> Self classmethod ¤

Initialize a PengIndependentAtomPotential with a convenience wrapper for the scattering factor parameters.

Arguments:

  • atom_positions: The coordinates of the atoms in units of angstroms.
  • parameters: A pytree for the scattering factor parameters from Peng et al. (1996).
  • extra_b_factors: Additional per-atom B-factors that are added to the values in scattering_parameters.b.
cryojax.simulator.PengIndependentAtomPotential.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> Self ¤

Since this class is itself an implementation of an AbstractVolumeParametrisation, this function maps to the identity.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.
cryojax.simulator.PengIndependentAtomPotential.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self ¤

Return a new potential with rotated atom_positions.

cryojax.simulator.PengIndependentAtomPotential.translate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with rotated atom_positions.

cryojax.simulator.PengIndependentAtomPotential.to_real_voxel_grid(shape: tuple[int, int, int], voxel_size: Float[NDArrayLike, ''] | float, *, batch_options: dict[str, Any] = {}) -> Float[Array, '{shape[0]} {shape[1]} {shape[2]}'] ¤

Return a voxel grid of the potential in real space.

Through the work of Peng et al. (1996), tabulated elastic electron scattering factors are defined as

\[f^{(e)}(\mathbf{q}) = \sum\limits_{i = 1}^5 a_i \exp(- b_i |\mathbf{q}|^2),\]

where \(a_i\) is stored as PengAtomicPotential.scattering_factor_a and \(b_i\) is stored as PengAtomicPotential.scattering_factor_b for the scattering vector \(\mathbf{q}\). Under usual scattering approximations (i.e. the first-born approximation), the rescaled electrostatic potential energy \(U(\mathbf{r})\) is then given by \(4 \pi \mathcal{F}^{-1}[f^{(e)}(\boldsymbol{\xi} / 2)](\mathbf{r})\), which is computed analytically as

\[U(\mathbf{r}) = 4 \pi \sum\limits_{i = 1}^5 \frac{a_i}{(2\pi (b_i / 8 \pi^2))^{3/2}} \exp(- \frac{|\mathbf{r} - \mathbf{r}'|^2}{2 (b_i / 8 \pi^2)}),\]

where \(\mathbf{r}'\) is the position of the atom. Including an additional B-factor (denoted by \(B\)) gives the expression for the potential \(U(\mathbf{r})\) of a single atom type and its fourier transform pair \(\tilde{U}(\boldsymbol{\xi}) \equiv \mathcal{F}[U](\boldsymbol{\xi})\),

\[U(\mathbf{r}) = 4 \pi \sum\limits_{i = 1}^5 \frac{a_i}{(2\pi ((b_i + B) / 8 \pi^2))^{3/2}} \exp(- \frac{|\mathbf{r} - \mathbf{r}'|^2}{2 ((b_i + B) / 8 \pi^2)}),\]
\[\tilde{U}(\boldsymbol{\xi}) = 4 \pi \sum\limits_{i = 1}^5 a_i \exp(- (b_i + B) |\boldsymbol{\xi}|^2 / 4) \exp(2 \pi i \boldsymbol{\xi}\cdot\mathbf{r}'),\]

where \(\mathbf{q} = \boldsymbol{\xi} / 2\) gives the relationship between the wave vector and the scattering vector.

In practice, for a discretization on a grid with voxel size \(\Delta r\) and grid point \(\mathbf{r}_{\ell}\), the potential is evaluated as the average value inside the voxel

\[U_{\ell} = 4 \pi \frac{1}{\Delta r^3} \sum\limits_{i = 1}^5 a_i \prod\limits_{j = 1}^3 \int_{r^{\ell}_j-\Delta r/2}^{r^{\ell}_j+\Delta r/2} dr_j \ \frac{1}{{\sqrt{2\pi ((b_i + B) / 8 \pi^2)}}} \exp(- \frac{(r_j - r'_j)^2}{2 ((b_i + B) / 8 \pi^2)}),\]

where \(j\) indexes the components of the spatial coordinate vector \(\mathbf{r}\). The above expression is evaluated using the error function as

\[U_{\ell} = 4 \pi \frac{1}{(2 \Delta r)^3} \sum\limits_{i = 1}^5 a_i \prod\limits_{j = 1}^3 \textrm{erf}(\frac{r_j^{\ell} - r'_j + \Delta r / 2}{\sqrt{2 ((b_i + B) / 8\pi^2)}}) - \textrm{erf}(\frac{r_j^{\ell} - r'_j - \Delta r / 2}{\sqrt{2 ((b_i + B) / 8\pi^2)}}).\]

Arguments:

  • shape: The shape of the resulting voxel grid.
  • voxel_size: The voxel size of the resulting voxel grid.
  • batch_options: Advanced options for controlling batching. This is a dictionary with the following keys:
    • "batch_size": The number of z-planes to evaluate in parallel with jax.vmap. By default, 1.
    • "n_batches": The number of iterations used to evaluate the volume, where the iteration is taken over groups of atoms. This is useful if batch_size = 1 and GPU memory is exhausted. By default, 1.

Returns:

The rescaled potential \(U_{\ell}\) as a voxel grid of shape shape and voxel size voxel_size.

Voxel-based¤

Fourier-space¤

Fourier-space conventions

  • The fourier_voxel_grid and frequency_slice arguments to FourierVoxelGridVolume.__init__ should be loaded with the zero frequency component in the center of the box. This is returned by the
  • The parameters in an AbstractPose represent a rotation in real-space. This means that when calling FourierVoxelGridVolume.rotate_to_pose, frequencies are rotated by the inverse rotation as stored in the pose.

cryojax.simulator.FourierVoxelGridVolume ¤

A 3D voxel grid in fourier-space.

cryojax.simulator.FourierVoxelGridVolume.__init__(fourier_voxel_grid: Complex[NDArrayLike, 'dim dim dim'], frequency_slice_in_pixels: Float[NDArrayLike, '1 dim dim 3']) ¤

Arguments:

  • fourier_voxel_grid: The cubic voxel grid in fourier space.
  • frequency_slice_in_pixels: The frequency slice coordinate system.
cryojax.simulator.FourierVoxelGridVolume.from_real_voxel_grid(real_voxel_grid: Float[NDArrayLike, 'dim dim dim'], *, pad_scale: float = 1.0, pad_mode: str = 'constant', filter: Optional[AbstractFilter] = None) -> Self classmethod ¤

Load from a real-valued 3D voxel grid.

Arguments:

  • real_voxel_grid: A voxel grid in real space.
  • pad_scale: Scale factor at which to pad real_voxel_grid before fourier transform. Must be a value greater than 1.0.
  • pad_mode: Padding method. See jax.numpy.pad for documentation.
  • filter: A filter to apply to the result of the fourier transform of real_voxel_grid, i.e. fftn(real_voxel_grid). Note that the zero frequency component is assumed to be in the corner.
cryojax.simulator.FourierVoxelGridVolume.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> Self ¤

Since this class is itself an implementation of an AbstractVolumeParametrisation, this function maps to the identity.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.
cryojax.simulator.FourierVoxelGridVolume.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self ¤

Return a new volume with a rotated frequency_slice_in_pixels.

cryojax.simulator.FourierVoxelGridVolume.frequency_slice_in_pixels: Float[Array, '1 dim dim 3'] = jnp.asarray(frequency_slice_in_pixels, dtype=float) instance-attribute ¤
cryojax.simulator.FourierVoxelGridVolume.shape: tuple[int, int, int] property ¤

The shape of the fourier_voxel_grid.

cryojax.simulator.FourierVoxelSplineVolume ¤

A 3D voxel grid in fourier-space, represented by spline coefficients.

cryojax.simulator.FourierVoxelSplineVolume.__init__(spline_coefficients: Complex[NDArrayLike, 'coeff_dim coeff_dim coeff_dim'], frequency_slice_in_pixels: Float[NDArrayLike, '1 dim dim 3']) ¤

Arguments:

  • spline_coefficients: The spline coefficents computed from the cubic voxel grid in fourier space. See cryojax.ndimage.compute_spline_coefficients.
  • frequency_slice_in_pixels: Frequency slice coordinate system. See cryojax.coordinates.make_frequency_slice.
cryojax.simulator.FourierVoxelSplineVolume.from_real_voxel_grid(real_voxel_grid: Float[NDArrayLike, 'dim dim dim'], *, pad_scale: float = 1.0, pad_mode: str = 'constant', filter: Optional[AbstractFilter] = None) -> Self classmethod ¤

Load from a real-valued 3D voxel grid.

Arguments:

  • real_voxel_grid: A voxel grid in real space.
  • pad_scale: Scale factor at which to pad real_voxel_grid before fourier transform. Must be a value greater than 1.0.
  • pad_mode: Padding method. See jax.numpy.pad for documentation.
  • filter: A filter to apply to the result of the fourier transform of real_voxel_grid, i.e. fftn(real_voxel_grid). Note that the zero frequency component is assumed to be in the corner.
cryojax.simulator.FourierVoxelSplineVolume.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> Self ¤

Since this class is itself an implementation of an AbstractVolumeParametrisation, this function maps to the identity.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.
cryojax.simulator.FourierVoxelSplineVolume.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self ¤

Return a new volume with a rotated frequency_slice_in_pixels.

cryojax.simulator.FourierVoxelSplineVolume.frequency_slice_in_pixels: Float[Array, '1 dim dim 3'] = jnp.asarray(frequency_slice_in_pixels, dtype=float) instance-attribute ¤
cryojax.simulator.FourierVoxelSplineVolume.shape: tuple[int, int, int] property ¤

The shape of the original fourier_voxel_grid from which coefficients were computed.

Real-space¤

cryojax.simulator.RealVoxelGridVolume ¤

A 3D voxel grid in real-space.

cryojax.simulator.RealVoxelGridVolume.__init__(real_voxel_grid: Float[NDArrayLike, 'dim dim dim'], coordinate_grid_in_pixels: Float[NDArrayLike, 'dim dim dim 3']) ¤

Arguments:

  • real_voxel_grid: The voxel grid in fourier space.
  • coordinate_grid_in_pixels: A coordinate grid.
cryojax.simulator.RealVoxelGridVolume.from_real_voxel_grid(real_voxel_grid: Float[NDArrayLike, 'dim dim dim'], *, coordinate_grid_in_pixels: Optional[Float[Array, 'dim dim dim 3']] = None, crop_scale: Optional[float] = None) -> Self classmethod ¤

Load a RealVoxelGridVolume from a real-valued 3D voxel grid.

Arguments:

  • real_voxel_grid: A voxel grid in real space.
  • crop_scale: Scale factor at which to crop real_voxel_grid. Must be a value greater than 1.
cryojax.simulator.RealVoxelGridVolume.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> Self ¤

Since this class is itself an implementation of an AbstractVolumeParametrisation, this function maps to the identity.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.
cryojax.simulator.RealVoxelGridVolume.rotate_to_pose(pose: AbstractPose, inverse: bool = False) -> Self ¤

Return a new volume with a rotated coordinate_grid_in_pixels.

cryojax.simulator.RealVoxelGridVolume.coordinate_grid_in_pixels: Float[Array, 'dim dim dim 3'] = jnp.asarray(coordinate_grid_in_pixels, dtype=float) instance-attribute ¤
cryojax.simulator.RealVoxelGridVolume.shape: tuple[int, int, int] property ¤

The shape of the real_voxel_grid.

Parametrising scattering potentials¤

cryojax.simulator.AbstractPotentialParametrisation

cryojax.simulator.AbstractPotentialParametrisation ¤

Abstract interface for a scattering potential.

Info

In, cryojax, potentials should be built in units of inverse length squared, \([L]^{-2}\). This rescaled potential is defined to be

\[U(\mathbf{r}) = \frac{2 m e}{\hbar^2} V(\mathbf{r}),\]

where \(V\) is the electrostatic potential energy, \(\mathbf{r}\) is a positional coordinate, \(m\) is the electron mass, and \(e\) is the electron charge.

For a single atom, this rescaled potential has the advantage that under usual scattering approximations (i.e. the first-born approximation), the fourier transform of this quantity is closely related to tabulated electron scattering factors. In particular, for a single atom with scattering factor \(f^{(e)}(\mathbf{q})\) and scattering vector \(\mathbf{q}\), its rescaled potential is equal to

\[U(\mathbf{r}) = 4 \pi \mathcal{F}^{-1}[f^{(e)}(\boldsymbol{\xi} / 2)](\mathbf{r}),\]

where \(\boldsymbol{\xi} = 2 \mathbf{q}\) is the wave vector coordinate and \(\mathcal{F}^{-1}\) is the inverse fourier transform operator in the convention

\[\mathcal{F}[f](\boldsymbol{\xi}) = \int d^3\mathbf{r} \ \exp(2\pi i \boldsymbol{\xi}\cdot\mathbf{r}) f(\mathbf{r}).\]

The rescaled potential \(U\) gives the following time-independent schrodinger equation for the scattering problem,

\[(\nabla^2 + k^2) \psi(\mathbf{r}) = - U(\mathbf{r}) \psi(\mathbf{r}),\]

where \(k\) is the incident wavenumber of the electron beam.

References:

  • For the definition of the rescaled potential, see Chapter 69, Page 2003, Equation 69.6 from Hawkes, Peter W., and Erwin Kasper. Principles of Electron Optics, Volume 4: Advanced Wave Optics. Academic Press, 2022.
  • To work out the correspondence between the rescaled potential and the electron scattering factors, see the supplementary information from Vulović, Miloš, et al. "Image formation modeling in cryo-electron microscopy." Journal of structural biology 183.1 (2013): 19-32.
cryojax.simulator.AbstractPotentialParametrisation.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation abstractmethod ¤

Core interface for computing the representation of the volume.

cryojax.simulator.AbstractTabulatedPotential

cryojax.simulator.AbstractTabulatedPotential ¤

Abstract class for a scattering potential from tabulated electron scattering factors.

cryojax.simulator.AbstractTabulatedPotential.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation abstractmethod ¤

Core interface for computing the representation of the volume.

cryojax.simulator.AbstractTabulatedPotential.from_tabulated_parameters(*args: Any, **kwargs: Any) -> Self abstractmethod classmethod ¤

Construct a scattering potential parametrisation from tabulated electron scattering factors.

cryojax.simulator.AbstractPengPotential

cryojax.simulator.AbstractPengPotential ¤

Abstract class for a scattering potential parametrised as a mixture of five gaussians per atom (Peng et al. 1996).

References:

  • Peng, L-M. "Electron atomic scattering factors and scattering potentials of crystals." Micron 30.6 (1999): 625-648.
  • Peng, L-M., et al. "Robust parameterization of elastic and absorptive electron atomic scattering factors." Acta Crystallographica Section A: Foundations of Crystallography 52.2 (1996): 257-276.
cryojax.simulator.AbstractPengPotential.amplitudes: eqx.AbstractVar[Float[Array, 'n_atoms n_gaussians']] instance-attribute ¤
cryojax.simulator.AbstractPengPotential.b_factors: eqx.AbstractVar[Float[Array, 'n_atoms n_gaussians']] instance-attribute ¤
cryojax.simulator.AbstractPengPotential.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation abstractmethod ¤

Core interface for computing the representation of the volume.

cryojax.simulator.AbstractPengPotential.from_tabulated_parameters(*args: Any, **kwargs: Any) -> Self abstractmethod classmethod ¤

Construct a scattering potential parametrisation from tabulated electron scattering factors.

Parametrising ensembles¤

cryojax.simulator.AbstractEnsembleParametrisation

cryojax.simulator.AbstractEnsembleParametrisation ¤

Abstract interface for a volume with conformational heterogeneity.

cryojax.simulator.AbstractEnsembleParametrisation.conformation: eqx.AbstractVar[Any] instance-attribute ¤

A variable for the ensemble's conformational state.

cryojax.simulator.AbstractEnsembleParametrisation.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation abstractmethod ¤

Core interface for computing the representation of the volume.

cryojax.simulator.DiscreteStructuralEnsemble ¤

Abstraction of an ensemble with discrete conformational heterogeneity.

cryojax.simulator.DiscreteStructuralEnsemble.__init__(conformational_space: tuple[AbstractVolumeRepresentation, ...], conformation: int | Int[Array, '']) ¤

Arguments:

  • conformational_space: A tuple of AbstractStructureRepresentations.
  • conformation: A conformation with a discrete index at which to evaluate the tuple.
cryojax.simulator.DiscreteStructuralEnsemble.compute_volume_representation(rng_key: Optional[PRNGKeyArray] = None) -> AbstractVolumeRepresentation ¤

Map to the volume at conformation.

Arguments:

  • rng_key: Not used in this implementation, but optionally included for other implementations.