Skip to content

Your entry point¤

cryojax.simulator.make_image_model ¤

make_image_model(
    volume: cryojax.simulator.AbstractVolumeParametrization,
    image_config: cryojax.simulator.AbstractImageConfig,
    pose: cryojax.simulator.AbstractPose,
    transfer_theory: cryojax.simulator.ContrastTransferTheory | None = None,
    volume_integrator: cryojax.simulator.AbstractVolumeIntegrator = cryojax.simulator.AutoVolumeProjection(),
    detector: cryojax.simulator.AbstractDetector | None = None,
    *,
    image_transform: cryojax.ndimage.AbstractImageTransform | None = None,
    normalizes_signal: bool = False,
    signal_region: Bool[NDArrayLike, "_ _"] | None = None,
    signal_centering: Literal["bg", "mean"] = "mean",
    translate_mode: Literal["fft", "atom", "none"] = "fft",
    quantity_mode: Literal["contrast", "intensity", "counts", "none"] = "none",
) -> cryojax.simulator.AbstractImageModel

Construct an cryojax.simulator.AbstractImageModel for most common use-cases.

Simulate an image

import cryojax.simulator as cxs

# Load modeling components
volume, image_config, pose, transfer_theory = ...
# Build image formation model
image_model = cxs.make_image_model(volume, image_config, pose, transfer_theory)
# Simulate!
image = image_model.simulate()

Main arguments:

Options:

  • image_transform: A cryojax.ndimage.AbstractImageTransform applied to the the output of image_model.simulate() as a postprocessing step.
  • normalizes_signal: Whether or not to normalize the output of image_model.simulate(). If True, see signal_centering for options.
  • signal_region: A boolean array that is 1 where there is signal, and 0 otherwise used to normalize the image. Must have shape equal to image_config.shape.
  • signal_centering: How to calculate the offset for normalization when normalizes_signal = True (and ignored if normalizes_signal = False). Options are
    • 'mean': Normalize the image to be mean 0 within signal_region. This normalizes the image to be a z-score.
    • 'bg': Subtract mean value at the image edges. This makes the image fade to a background with values equal to zero. Requires that image_config.padded_shape is large enough so that the signal sufficiently decays.
  • translate_mode: How to apply in-plane translation to the volume. Options are
    • 'fft': Apply phase shifts in the Fourier domain. This option is best for most use cases and is usually faster than the 'atom' option.
    • 'atom': Apply translation to atom positions before projection. This method is more numerically accurate than the 'fft' option, but it is only supported if the volume argument yields a cryojax.simulator.AbstractAtomVolume.
    • 'none': Do not apply the translation.
  • quantity_mode: The physical observable to simulate. Options are:

Returns:

An cryojax.simulator.AbstractImageModel. This has type:

cryojax.simulator.load_tabulated_volume ¤

load_tabulated_volume(
    path_or_mmdf: str | pathlib.Path | pandas.DataFrame,
    *,
    output_type: type[
        cryojax.simulator.GaussianFourierVolume
        | cryojax.simulator.GaussianMixtureVolume
    ] = cryojax.simulator.GaussianMixtureVolume,
    tabulation: Literal["peng"] = "peng",
    include_b_factors: bool = False,
    b_factor_fn: Callable[
        [cryojax.jax_util.NDArrayLike, cryojax.jax_util.NDArrayLike],
        cryojax.jax_util.NDArrayLike,
    ] = identity_fn,
    selection_string: str = "all",
    pdb_options: dict[str, Any] = {},
) -> (
    cryojax.simulator.GaussianFourierVolume | cryojax.simulator.GaussianMixtureVolume
)

Load an atomistic representation of a volume from tabulated electron scattering factors.

Warning

This function cannot be used with JIT compilation. Rather, its output should be passed to JIT-compiled functions. For example:

import cryojax.simulator as cxs
import equinox as eqx

path_to_pdb = ...
volume = cxs.load_tabulated_volume(path_to_pdb)

@eqx.filter_jit
def simulate_fn(volume, ...):
    image_model = cxs.make_image_model(volume, ...)
    return image_model.simulate()

image = simulate_fn(volume, ...)

Arguments:

Returns:

A cryojax.simulator.AbstractVoxelVolume with exact type equal to output_type.

cryojax.simulator.suggest_n_spread ¤

suggest_n_spread(
    volume: cryojax.simulator.GaussianMixtureVolume,
    pixel_size: cryojax.jax_util.FloatLike,
    *,
    cutoff_sigma: float = 4.0,
    mode: Literal["global", "termwise"] = "termwise",
    min_n_spread: int = 1,
    max_n_spread: int | None = None,
) -> int | tuple[int, ...]

Suggest n_spread value(s) for real-space gaussian spreading, sized to volume's variances via cryojax.ndimage.variance_to_nspread.

Usage

Call this on the concrete volume, before it enters jax.jit, and pass the result to the matching concrete integrator/render-fn -- not AutoVolumeProjection/AutoVolumeRenderFn, which can't resolve n_spread from a traced volume and don't save you anything here anyway (you already need to know the volume's type to call this):

import jax
import cryojax.simulator as cxs

volume = cxs.load_tabulated_volume(
    path_to_pdb, output_type=cxs.GaussianMixtureVolume
)
n_spread = cxs.suggest_n_spread(volume, pixel_size)
integrator = cxs.GaussianMixtureProjection(n_spread=n_spread)

image_model = jax.jit(make_image_model)(
    volume=volume, volume_integrator=integrator, ...
)

Arguments:

  • volume: The GaussianMixtureVolume to suggest n_spread for.
  • pixel_size: The pixel/voxel size of the grid volume is spread onto.
  • cutoff_sigma: Truncation width in standard deviations (n_sigma in variance_to_nspread).
  • mode: "termwise" (default) returns one n_spread per gaussian component, each sized to that component's largest variance across positions. "global" returns a single int sized to the single widest gaussian in volume.
  • min_n_spread/max_n_spread: Clamp every returned value to this range.

Returns:

An int (mode="global") or tuple[int, ...], one per gaussian component (mode="termwise").

cryojax.simulator.render_voxel_volume ¤

render_voxel_volume(
    atom_volume: cryojax.simulator.AbstractAtomVolume,
    render_fn: cryojax.simulator.AbstractVolumeRenderFn,
    *,
    output_type: type[
        cryojax.simulator.FourierVoxelGridVolume
        | cryojax.simulator.RealVoxelGridVolume
    ] = cryojax.simulator.FourierVoxelGridVolume,
    options: Mapping[str, Any] = {},
) -> (
    cryojax.simulator.FourierVoxelGridVolume | cryojax.simulator.RealVoxelGridVolume
)

Render a voxel volume representation from an atomistic one.

Simulate an image with Fourier slice extraction

import cryojax.simulator as cxs

# Simulate an image with Fourier slice extraction
voxel_volume = cxs.render_voxel_volume(
    atom_volume=cxs.load_tabulated_volume("example.pdb"),
    render_fn=cxs.AutoVolumeRenderFn(shape=(100, 100, 100), voxel_size=1.0),
    output_type=cxs.FourierVoxelGridVolume,
)
image_model = cxs.make_image_model(voxel_volume, ...)
image = image_model.simulate()

Arguments:

Returns:

A cryojax.simulator.AbstractVoxelVolume with exact type equal to output_type.