Simulate particle stacks

cryospax.simulate_particle_stack(dataset: cryospax.AbstractParticleDataset, simulate_fn: Callable[[PyTree, ~ConstantT, ~PerParticleT], Float[Array, '_ _']], constant_args: ~ConstantT = None, per_particle_args: ~PerParticleT = None, batch_size: int | None = None, images_per_file: int | None = None, **kwargs: Any)

Write a stack of images from parameters contained in a cryospax.AbstractParticleDataset.

Note

This function works generally for a simulate_fn of the form

image = simulate_fn(
    parameter_info, constant_args, per_particle_args
)

where parameter_info is the pytree read from the AbstractParticleDataset.parameter_file, constant_args is a parameter that does not change between images, and per_particle_args is a pytree whose leaves have a batch dimension equal to the number of particles to be simulated.

Example 1: Basic usage such as instantiating an cryospax.AbstractParticleDataset and writing a simulate_fn

import cryojax.simulator as cxs
import cryospax as spx
import jax
from jaxtyping import PyTree

# Load a `RelionParticleDataset` object. This loads
# parameters and writes images
dataset = spx.RelionParticleDataset(..., mode='w')

# Write your `simulate_fn` function, building an
# `AbstractImageModel` (see tutorials for details)

def simulate_fn(
    parameter_info: PyTree, # loaded from `dataset.parameter_file`
    constant_args: PyTree,
    _,
) -> jax.Array:
    # `constant_args` do not change between images. For
    # example, include the method of taking projections
    ... = constant_args
    # Using the pose, CTF, and image config from the
    # `parameter_info`, build image simulation model
    image_model = cxs.make_image_model(...)
    # ... and compute
    return image_model.simulate()

# Simulate images and write to disk
spx.simulate_particle_stack(
    dataset,
    simulate_fn,
    constant_args=(...)
    per_particle_args=None, # default
    batch_size=10,
)

Example 2: More-advanced usage, writing a simulate_fn that simulates images with noise. Uses per_particle_args as well as constant_args.

import cryojax.simulator as cxs
import cryospax as spx
import jax
from jaxtyping import Array, PyTree, Shaped

# Load a `RelionParticleDataset` object. This loads
# parameters and writes images
dataset = spx.RelionParticleDataset(..., mode='w')

# Instantiate per-particle arguments. First, the RNG keys used
# to generate the noise
seed = 0
key = jax.random.key(seed)
key, *keys_noise = jax.random.split(key, n_images+1)
keys_noise = jnp.array(keys_noise)
# ... then, add a scaling parameter for the images
key, subkey = jax.random.split(key)
scaling_params = jax.random.uniform(subkey, shape=(n_images,))

# Now write your `simulate_fn` function, building a
# `cryojax.simulator.GaussianWhiteNoiseModel` to
# simulate images with white noise (see tutorials for details)

def simulate_fn(
    parameter_info: PyTree,
    constant_args: PyTree,
    per_particle_args: PyTree[Shaped[Array, "_ ..."]],
) -> jax.Array:
    ... = constant_args
    key, scale = per_particle_args

    # Combine two previously split PyTrees
    image_model = cxs.make_image_model(...)
    distribution = cxs.GaussianWhiteNoiseModel(image_model, ...)

    return scale * distribution.sample(key)

spx.simulate_particle_stack(
    dataset,
    simulate_fn,
    constant_args=(...)
    per_particle_args=(keys_noise, scaling_params)
    batch_size=10,
)

Arguments:

• dataset: The cryospax.AbstractParticleDataset class. Note that this must be passed in writing mode, i.e. mode = 'w', and the parameter_file must be instantiated with options=dict(loads_metadata=False, ...), which is the default.
• simulate_fn: A callable that computes the image stack from the parameters contained in the STAR file.
• constant_args: The constant arguments to pass to the simulate_fn function. These must be the same for all images.
• per_particle_args: Arguments to pass to the simulate_fn function. This is a pytree with leaves having a batch size with equal dimension to the number of images.
• batch_size: The number images to compute in parallel using jax.vmap. If None, simulate images in a python for-loop. This is useful if the user isn't yet familiar with debugging JIT compilation.
• images_per_file: The number of images to write in a single image file. By default, set this as the number of particles in the dataset.
• kwargs: Keyword arguments passed to cryospax.AbstractParticleParameterFile.save