Skip to content

Scattering potential representations¤

cryojax provides different options for how to represent spatial potential energy distributions in cryo-EM.

cryojax.simulator.AbstractPotentialRepresentation

cryojax.simulator.AbstractPotentialRepresentation ¤

Abstract interface for the spatial potential energy distribution of a scatterer.

cryojax.simulator.AbstractPotentialRepresentation.rotate_to_pose(pose: AbstractPose) -> Self abstractmethod ¤

Return a new AbstractPotentialRepresentation at the given pose.

Arguments:

  • pose: The pose at which to view the AbstractPotentialRepresentation.

Atom-based scattering potentials¤

cryojax.simulator.AbstractAtomicPotential

cryojax.simulator.AbstractAtomicPotential ¤

Abstract interface for an atom-based scattering potential representation.

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.AbstractAtomicPotential.atom_positions: eqx.AbstractVar[Float[Array, 'n_atoms 3']] instance-attribute ¤
cryojax.simulator.AbstractAtomicPotential.as_real_voxel_grid(shape: tuple[int, int, int], voxel_size: Float[Array, ''] | float) -> Float[Array, '{shape[0]} {shape[1]} {shape[2]}'] abstractmethod ¤

cryojax.simulator.GaussianMixtureAtomicPotential ¤

An atomistic representation of scattering potential as a mixture of gaussians.

The naming and numerical convention of parameters gaussian_amplitudes and gaussian_widths follows "Robust Parameterization of Elastic and Absorptive Electron Atomic Scattering Factors" by Peng et al. (1996), where \(a_i\) are the gaussian_amplitudes and \(b_i\) are the gaussian_widths.

Info

In order to load a GaussianMixtureAtomicPotential from tabulated scattering factors, use the cryojax.constants submodule.

from cryojax.constants import (
    peng_element_scattering_factor_parameter_table,
    get_tabulated_scattering_factor_parameters,
)
from cryojax.io import read_atoms_from_pdb
from cryojax.simulator import GaussianMixtureAtomicPotential

# Load positions of atoms and one-hot encoded atom names
atom_positions, atom_identities = read_atoms_from_pdb(...)
scattering_factor_a, scattering_factor_b = get_tabulated_scattering_factor_parameters(
    atom_identities, peng_element_scattering_factor_parameter_table
)
potential = GaussianMixtureAtomicPotential(
    atom_positions, scattering_factor_a, scattering_factor_b
)
cryojax.simulator.GaussianMixtureAtomicPotential.__init__(atom_positions: Float[Array, 'n_atoms 3'] | Float[np.ndarray, 'n_atoms 3'], gaussian_amplitudes: Float[Array, 'n_atoms n_gaussians_per_atom'] | Float[np.ndarray, 'n_atoms n_gaussians_per_atom'], gaussian_widths: Float[Array, 'n_atoms n_gaussians_per_atom'] | Float[np.ndarray, 'n_atoms n_gaussians_per_atom']) ¤

Arguments:

  • atom_positions: The coordinates of the atoms in units of angstroms.
  • gaussian_amplitudes: The strength for each atom and for each gaussian per atom. This has units of angstroms.
  • gaussian_widths: The variance (up to numerical constants) for each atom and for each gaussian per atom. This has units of angstroms squared.
cryojax.simulator.GaussianMixtureAtomicPotential.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with rotated atom_positions.

cryojax.simulator.GaussianMixtureAtomicPotential.as_real_voxel_grid(shape: tuple[int, int, int], voxel_size: Float[Array, ''] | float, *, z_planes_in_parallel: int = 1, atom_groups_in_series: int = 1) -> Float[Array, '{shape[0]} {shape[1]} {shape[2]}'] ¤

Return a voxel grid of the potential in real space.

See PengAtomicPotential.as_real_voxel_grid for the numerical conventions used when computing the sum of gaussians.

Arguments:

  • shape: The shape of the resulting voxel grid.
  • voxel_size: The voxel size of the resulting voxel grid.
  • z_planes_in_parallel: The number of z-planes to evaluate in parallel with jax.vmap. By default, 1.
  • atom_groups_in_series: The number of iterations used to evaluate the volume, where the iteration is taken over groups of atoms. This is useful if z_planes_in_parallel = 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.

cryojax.simulator.PengAtomicPotential ¤

The scattering potential parameterized as a mixture of five gaussians per atom, through work by Lian-Mao Peng.

To load this object, the following pattern can be used:

from cryojax.io import read_atoms_from_pdb
from cryojax.simulator import PengAtomicPotential

# Load positions of atoms and one-hot encoded atom names
filename = "example.pdb"
atom_positions, atom_identities = read_atoms_from_pdb(filename)
potential = PengAtomicPotential(atom_positions, atom_identities)

Alternatively, use the following to load with B-factors:

from cryojax.io import read_atoms_from_pdb
from cryojax.simulator import PengAtomicPotential

# Load positions of atoms, encoded atom names, and B-factors
filename = "example.pdb"
atom_positions, atom_identities, b_factors = read_atoms_from_pdb(
    filename, get_b_factors=True
)
potential = PengAtomicPotential(atom_positions, atom_identities, b_factors)

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.PengAtomicPotential.__init__(atom_positions: Float[Array, 'n_atoms 3'] | Float[np.ndarray, 'n_atoms 3'], atom_identities: Int[Array, ' n_atoms'] | Int[np.ndarray, ' n_atoms'], b_factors: Optional[Float[Array, ' n_atoms'] | Float[np.ndarray, ' n_atoms']] = None) ¤

Arguments:

  • atom_positions: The coordinates of the atoms in units of angstroms.
  • atom_identities: Array containing the index of the one-hot encoded atom names. Hydrogen is "1", Carbon is "6", Nitrogen is "7", etc.
  • b_factors: The B-factors applied to each atom.
cryojax.simulator.PengAtomicPotential.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with rotated atom_positions.

cryojax.simulator.PengAtomicPotential.as_real_voxel_grid(shape: tuple[int, int, int], voxel_size: Float[Array, ''] | float, *, z_planes_in_parallel: int = 1, atom_groups_in_series: int = 1) -> 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\) and stored as PengAtomicPotential.b_factors) 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.
  • z_planes_in_parallel: The number of z-planes to evaluate in parallel with jax.vmap. By default, 1.
  • atom_groups_in_series: The number of iterations used to evaluate the volume, where the iteration is taken over groups of atoms. This is useful if z_planes_in_parallel = 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 scattering potentials¤

cryojax.simulator.AbstractVoxelPotential

cryojax.simulator.AbstractVoxelPotential ¤

Abstract interface for a voxel-based scattering potential representation.

cryojax.simulator.AbstractVoxelPotential.voxel_size: AbstractVar[Float[Array, '']] instance-attribute ¤
cryojax.simulator.AbstractVoxelPotential.is_real: AbstractClassVar[bool] instance-attribute ¤
cryojax.simulator.AbstractVoxelPotential.shape: tuple[int, ...] abstractmethod property ¤

The shape of the voxel array.

cryojax.simulator.AbstractVoxelPotential.from_real_voxel_grid(real_voxel_grid: Float[Array, 'dim dim dim'] | Float[np.ndarray, 'dim dim dim'], voxel_size: Float[Array, ''] | Float[np.ndarray, ''] | float) -> Self abstractmethod classmethod ¤

Load an AbstractVoxelPotential from real-valued 3D electron scattering potential.

Fourier-space voxel representations¤

Fourier-space conventions

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

cryojax.simulator.FourierVoxelGridPotential ¤

A 3D scattering potential voxel grid in fourier-space.

cryojax.simulator.FourierVoxelGridPotential.__init__(fourier_voxel_grid: Complex[Array, 'dim dim dim'], frequency_slice_in_pixels: Float[Array, '1 dim dim 3'], voxel_size: Float[Array, ''] | float) ¤

Arguments:

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

Load an AbstractFourierVoxelGridPotential from real-valued 3D electron scattering potential voxel grid.

Arguments:

  • real_voxel_grid: A scattering potential voxel grid in real space.
  • voxel_size: The voxel size of real_voxel_grid.
  • 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.FourierVoxelGridPotential.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with a rotated frequency_slice_in_pixels.

cryojax.simulator.FourierVoxelGridPotential.frequency_slice_in_angstroms: Float[Array, '1 dim dim 3'] cached property ¤

The frequency_slice_in_pixels in angstroms.

cryojax.simulator.FourierVoxelGridPotential.shape: tuple[int, int, int] property ¤

The shape of the fourier_voxel_grid.

cryojax.simulator.FourierVoxelGridPotentialInterpolator ¤

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

cryojax.simulator.FourierVoxelGridPotentialInterpolator.__init__(fourier_voxel_grid: Float[Array, 'dim dim dim'], frequency_slice_in_pixels: Float[Array, '1 dim dim 3'], voxel_size: Float[Array, ''] | float) ¤

Note

The argument fourier_voxel_grid is used to set FourierVoxelGridPotentialInterpolator.coefficients in the __init__, but it is not stored in the class. For example,

voxels = FourierVoxelGridPotentialInterpolator(
    fourier_voxel_grid, frequency_slice, voxel_size
)
assert hasattr(voxels, "fourier_voxel_grid")  # This will return an error
assert hasattr(voxels, "coefficients")  # Instead, store spline coefficients

Arguments:

  • fourier_voxel_grid: The cubic voxel grid in fourier space.
  • frequency_slice_in_pixels: Frequency slice coordinate system, wrapped in a FrequencySlice object.
  • voxel_size: The voxel size.
cryojax.simulator.FourierVoxelGridPotentialInterpolator.from_real_voxel_grid(real_voxel_grid: Float[Array, 'dim dim dim'] | Float[np.ndarray, 'dim dim dim'], voxel_size: Float[Array, ''] | Float[np.ndarray, ''] | float, *, pad_scale: float = 1.0, pad_mode: str = 'constant', filter: Optional[AbstractFilter] = None) -> Self classmethod ¤

Load an AbstractFourierVoxelGridPotential from real-valued 3D electron scattering potential voxel grid.

Arguments:

  • real_voxel_grid: A scattering potential voxel grid in real space.
  • voxel_size: The voxel size of real_voxel_grid.
  • 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.FourierVoxelGridPotentialInterpolator.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with a rotated frequency_slice_in_pixels.

cryojax.simulator.FourierVoxelGridPotentialInterpolator.frequency_slice_in_angstroms: Float[Array, '1 dim dim 3'] cached property ¤

The frequency_slice_in_pixels in angstroms.

cryojax.simulator.FourierVoxelGridPotentialInterpolator.shape: tuple[int, int, int] property ¤

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

Real-space voxel representations¤

cryojax.simulator.RealVoxelGridPotential ¤

Abstraction of a 3D scattering potential voxel grid in real-space.

cryojax.simulator.RealVoxelGridPotential.__init__(real_voxel_grid: Float[Array, 'dim dim dim'], coordinate_grid_in_pixels: Float[Array, 'dim dim dim 3'], voxel_size: Float[Array, ''] | float) ¤

Arguments:

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

Load a RealVoxelGridPotential from a real-valued 3D electron scattering potential voxel grid.

Arguments:

  • real_voxel_grid: An electron scattering potential voxel grid in real space.
  • voxel_size: The voxel size of real_voxel_grid.
  • crop_scale: Scale factor at which to crop real_voxel_grid. Must be a value greater than 1.
cryojax.simulator.RealVoxelGridPotential.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with a rotated coordinate_grid_in_pixels.

cryojax.simulator.RealVoxelGridPotential.coordinate_grid_in_angstroms: Float[Array, 'dim dim dim 3'] cached property ¤

The coordinate_grid_in_pixels in angstroms.

cryojax.simulator.RealVoxelGridPotential.shape: tuple[int, int, int] property ¤

The shape of the real_voxel_grid.

cryojax.simulator.RealVoxelCloudPotential ¤

Abstraction of a 3D electron scattering potential voxel point cloud.

Info

This object is similar to the RealVoxelGridPotential. Instead of storing the whole voxel grid, a RealVoxelCloudPotential need only store points of non-zero scattering potential. Therefore, a RealVoxelCloudPotential stores a point cloud of scattering potential voxel values. Instantiating with the from_real_voxel_grid constructor will automatically mask points of zero scattering potential.

cryojax.simulator.RealVoxelCloudPotential.__init__(voxel_weights: Float[Array, ' size'], coordinate_list_in_pixels: Float[Array, 'size 3'], voxel_size: Float[Array, ''] | float) ¤

Arguments:

  • voxel_weights: A point-cloud of voxel scattering potential values.
  • coordinate_list_in_pixels: Coordinate list for the voxel_weights.
  • voxel_size: The voxel size.
cryojax.simulator.RealVoxelCloudPotential.from_real_voxel_grid(real_voxel_grid: Float[Array, 'dim dim dim'] | Float[np.ndarray, 'dim dim dim'], voxel_size: Float[Array, ''] | Float[np.ndarray, ''] | float, *, coordinate_grid_in_pixels: Optional[Float[Array, 'dim dim dim 3']] = None, rtol: float = 1e-05, atol: float = 1e-08, size: Optional[int] = None, fill_value: Optional[float] = None) -> Self classmethod ¤

Load an RealVoxelCloudPotential from a real-valued 3D electron scattering potential voxel grid.

Arguments:

  • real_voxel_grid: An electron scattering potential voxel grid in real space.
  • voxel_size: The voxel size of real_voxel_grid.
  • rtol: Argument passed to jnp.isclose, used for masking voxels of zero scattering potential.
  • atol: Argument passed to jnp.isclose, used for masking voxels of zero scattering potential.
  • size: Argument passed to jnp.where, used for fixing the size of the masked scattering potential. This argument is required for using this function with a JAX transformation.
  • fill_value: Argument passed to jnp.where, used if size is specified and the mask has fewer than the indicated number of elements.
cryojax.simulator.RealVoxelCloudPotential.rotate_to_pose(pose: AbstractPose) -> Self ¤

Return a new potential with a rotated coordinate_list_in_pixels.

cryojax.simulator.RealVoxelCloudPotential.coordinate_list_in_angstroms: Float[Array, 'size 3'] cached property ¤

The coordinate_list_in_pixels in angstroms.

cryojax.simulator.RealVoxelCloudPotential.shape: tuple[int] property ¤

The shape of voxel_weights.