Cryo-EM dataset manipulation

CryoSPAX implements interfaces for reading/writing to common cryo-EM software frameworks, such as RELION and cryoSPARC.

cryospax.AbstractDataset

cryospax.AbstractDataset

An abstraction of a dataset in cryospax. To create an AbstractDataset, implement its __init__, __getitem__, and __len__ methods.

This follows the torch.utils.data.Dataset API and can easily be wrapped into a pytorch Dataset with something like the following pattern:

import numpy as np
from torch.utils.data import Dataset

class CustomTorchDataset(Dataset):

    def __init__(cryojax_dataset: AbstractDataset):
        self.cryojax_dataset = cryojax_dataset

    def __getitem___(self, index) -> dict[str, np.ndarray]:
        particle_info = self.cryojax_dataset[index]
        return dict(index=index, images=np.asarray(particle_info["images"]))

    def __len__(self) -> int:
        return len(self.cryojax_dataset)

JAX also includes packages for dataloaders, such as jax-dataloaders and grain.

How do I implement an AbstractDataset?

Implementing an AbstractDataset is not like implementing classes in cryojax, which are equinox.Modules. An equinox.Module is just a pytree, so it can be safely passed to jax transformations. However, an AbstractDataset can not be passed to jax transformations. It is a normal python class, not a pytree.

__len__() -> int

__getitem__(index) -> ~T1

---

Datasets output a pytree that can be passed the JAX transformations. A particular challenge is passing these pytrees to jax.vmap; different cryoSPAX datasets may load different pytrees, and each pytree may have different arrays which are broadcasted.

cryospax.get_in_axes(dataset: cryospax.AbstractDataset) -> PyTree

For a function that accepts an output like value = dataset[0:10], this function returns a prefix pytree for input to equinox.filter_vmap.

Example

import cryospax as spx

dataset = spx.RelionParticleDataset(...)

@eqx.filter_vmap(in_axes=(spx.get_in_axes(dataset), None))
def fn_vmap(particle_info, args):
    ...

particle_info = dataset[0:10]
args = ...  # other arguments to `fn_vmap`
out = fn_vmap(particle_info, args)

Arguments:

• dataset: A cryospax.AbstractDataset.

Returns:

The in_axes argument for equinox.filter_vmap for the output of dataset[...].

Reading/writing parameter files (e.g. STAR files)

cryospax.AbstractParticleParameterFile

cryospax.AbstractParticleParameterFile(cryospax.AbstractDataset)

Base class for a particle parameter file.

path_to_output property

mode property

__getitem__(index) -> ~T1

__setitem__(index, value: ~T2)

__len__() -> int

append(value: ~T2)

save()

cryospax.AbstractRelionParticleParameterFile

cryospax.AbstractRelionParticleParameterFile(cryospax.AbstractParticleParameterFile)

Abstract class for a RELION particle parameter file.

This class is mostly useful for wrapping the cryospax.RelionParticleParameterFile into slightly different indexing behavior downstream.

__getitem__(index) -> ~T1

__setitem__(index, value: ~T2)

__len__() -> int

append(value: ~T2)

save()

path_to_starfile property

particle_data property

optics_data property

max_optics_groups property

loads_metadata property

loads_envelope property

make_image_config property

cryospax.RelionParticleParameterFile(cryospax.AbstractRelionParticleParameterFile)

A dataset that wraps a RELION particle stack in STAR format.

Matching cryoJAX and RELION rotation conventions

When simulating images with cryojax.simulator.make_image_model, conventions will always match between cryoJAX and RELION.

However, there are cases where it is necessary to correct for the difference with a call to pose.to_inverse_rotation(), where pose is the cryoJAX pose class loaded from the RelionParticleParameterFile.

Whether or not this correction should be made depends on if the pose represents a rotation of the object or the frame. This is encoded by the cryoJAX volume representation metadata volume.rotation_convention. When this is equal to 'frame' and not using make_image_model, call pose.to_inverse_rotation() to match with RELION.

__init__(path_to_starfile: str | pathlib.Path, mode: Literal['r', 'w'] = 'r', *, selection_filter: dict[str, Callable] = {}, exist_ok: bool = False, num_particles: int = 0, max_optics_groups: int | None = None, options: dict[str, Any] = {})

Arguments:

• path_to_starfile: The path to the RELION STAR file. If the path does not exist and mode = 'w', an empty dataset will be created.
• mode:
  • If mode = 'w', the dataset is prepared to write new parameters. This is done by storing an empty dataset in RelionParticleParameterFile.starfile_data. If a STAR file already exists at path_to_starfile, set exist_ok = True.
  • If mode = 'r', the STAR file at path_to_starfile is read into RelionParticleParameterFile.starfile_data.
• selection_filter: A dictionary used to include only particular dataset elements. The keys of this dictionary should be any data entry in the STAR file, while the values should be a function that takes in a column and returns a boolean mask for the column. For example, filter by class using selection_filter["rlnClassNumber"] = lambda x: x == 0.
• exist_ok: If the path_to_starfile already exists, if True and mode = 'w' nonetheless stores an empty RelionParticleParameterFile.starfile_data.
• num_particles: If in mode = 'w', initialize STAR file data to be num_particles entries. These entries are filled with NaN values and must be set via parameter_file[index] = parameter_info syntax.

• max_optics_groups: The maximum allowed optics group entries in the STAR file. The default value of this depends on if mode = 'r' or mode = 'w':

  • If mode = 'r': By default, max_optics_groups is twice the number of optics entries in the STAR file.
  • If mode = 'w': By default, max_optics_groups is equal to 1.

  Info

  This argument can be thought of as the number of allowed calls to parameter_file[...] = parameter_info or parameter_file.append(parameter_info) before an error will be thrown. Set this to a large value to be safe.

• options: A dictionary of options for modifying the behavior of reading/writing.

  • 'loads_metadata': If True, the resulting dict loads the raw metadata from the STAR file that is not otherwise included into a pandas.DataFrame. If this is set to True, note that dictionaries cannot pass through JIT boundaries without removing the metadata. By default, False.
  • 'loads_envelope': If True, read in the parameters of the CTF envelope function, i.e. "rlnCtfScalefactor" and "rlnCtfBfactor". By default, False.
  • 'make_image_config': A function with signature fn(shape, pixel_size, voltage_in_kilovolts) that returns a cryojax.simulator.BasicImageConfig class. Use this argument when it is desired to customize the image_config returned from this class, i.e. value = parameter_file[0:7]; print(value["image_config"]).

load(path_to_starfile: str | pathlib.Path, *, selection_filter: dict[str, Callable] = {}, max_optics_groups: int | None = None, loads_metadata: bool = False, loads_envelope: bool = False, make_image_config: Callable[[tuple[int, int], FloatLike, FloatLike], cryojax.simulator.BasicImageConfig] = default_fn) -> typing.Self classmethod

Convenience wrapper for cryospax.RelionParticleParameterFile.__init__ in mode = 'r' and reading via parameter_info = parameter_file[index] or writing via parameter_file[index] = parameter_info.

empty(path_to_starfile: str | pathlib.Path, num_particles: int, *, max_optics_groups: int = 1, exist_ok: bool = False, loads_envelope: bool = False) -> typing.Self classmethod

Convenience wrapper for cryospax.RelionParticleParameterFile.__init__ in mode = 'w' and writing via parameter_file[index] = parameter_info or parameter_file.append(parameter_info).

__getitem__(index: int | slice | Int[ndarray, ''] | Int[ndarray, '_']) -> dict[str, Any]

Load STAR file entries with value = parameter_file[...] syntax, where value is a dictionary with keys:

• 'pose': The cryojax.simulator.EulerAnglePose
• 'image_config': The cryojax.simulator.BasicImageConfig
• 'transfer_theory': The cryojax.simulator.ContrastTransferTheory
• 'metadata': If loads_metadata = True, a pandas.DataFrame of entries not used when loading the pose, image_config, and transfer_theory (e.g. the 'rlnClassNumber'). Otherwise, this key is not included.

__setitem__(index: int | slice | Int[ndarray, ''] | Int[ndarray, '_'], value: dict[str, Any])

Set STAR file entries with parameter_file[...] = value syntax, where value is a dictionary with keys:

• 'pose': The cryojax.simulator.EulerAnglePose
• 'image_config': The cryojax.simulator.BasicImageConfig
• 'transfer_theory': The cryojax.simulator.ContrastTransferTheory
• 'metadata': A pandas.DataFrame used to write custom entries to the STAR file. len(metadata) should be equal to the number of particles, and metadata.columns should be a subset of columns that already exist in parameter_file.particle_data but do not include entries filled by the image_config, transfer_theory, and pose.

All keys are optional, and unless image_config is not given, parameter_file[...] = value creates a new optics group.

Warning

Setting parameter_file[...] = value where value has a transfer_theory but not an image_config may result in surprising behavior: a new optics group will not be created, even though 'rlnSphericalAberration' and 'rlnAmplitudeContrast' are parameters stored in the transfer_theory.

This is because there is not enough information to create a new optics group if both are not given, and an error is not thrown so that patterns like

value = {"transfer_theory": cxs.ContrastTransferTheory(...), "pose": cxs.EulerAnglePose}(...)
parameter_file[index] = value

can be used to set particle entries.

__len__() -> int

The number of particles in the STAR file. This is simply an alias to cryospax.RelionParticleParameterFile.num_particles.

append(value: dict[str, Any])

Add an entry or entries to the STAR file with parameter_file.append(value) syntax, where value is a dictionary with keys:

• 'pose': The cryojax.simulator.EulerAnglePose This key is required.
• 'image_config': The cryojax.simulator.BasicImageConfig This key is required.
• 'transfer_theory': The cryojax.simulator.ContrastTransferTheory This key is required.
• 'metadata': A pandas.DataFrame used to write custom entries to the STAR file. len(metadata) should be equal to the number of particles, and metadata.columns should be a subset of columns that already exist in parameter_file.particle_data but do not include entries filled by the image_config, transfer_theory, and pose. This key is optional.

A new optics group will be created on each call to parameter_file.append(value).

copy() -> typing.Self

Make a deep copy of the parameter_file. This is often useful when modifying via parameter_file[...] = value syntax and it is necessary to maintain a copy of the original data.

save(*, overwrite: bool = False, **kwargs: Any)

Save the STAR file at the current parameter_file.path_to_starfile.

particle_data property

The pandas.DataFrame of particle STAR file entries.

optics_data property

The pandas.DataFrame of optics STAR file entries.

path_to_starfile property

The path to the STAR file. This is an alias to parameter_file.path_to_output.

Modify and write a new STAR file

parameter_file = RelionParticleParameterFile(
    path_to_starfile="./path/to/particles.star", mode="r"
)
updated_parameters = ...
parameter_file[...] = updated_parameters
parameter_file.path_to_starfile = "./path/to/new/particles.star"
parameter_file.save()

path_to_output property

mode property

Whether or not the parameter_file was instantiated in reading ('r') or writing ('w') mode.

This cannot be modified after initialization.

num_particles property

The number of particles in the STAR file.

num_optics_groups property

The number of optics groups in the STAR file.

max_optics_groups property

The maximum number of optics groups allowed in the STAR file.

Info

Every time parameters are written via parameter_file[...] = parameter_info or parameter_file.append(parameter_info) a new optics group is created.

Under the hood, the underlying STAR file data pre-allocates max_optics_groups entries for these writes and throws away any extra on calls to parameter_file.save().

loads_metadata property

Whether or not to load the following information:

parameter_info = parameter_file[index]
assert "metadata" in parameter_info  # True

The metadata is a pandas.DataFrame that includes the the STAR file rows not loaded by the parameter_file. For example, 'rlnClassNumber'.

loads_envelope property

Whether or not to load the following information:

parameter_info = parameter_file[index]
assert parameter_info["transfer_theory"].envelope is not None  # True

make_image_config property

A function that returns a cryojax.simulator.BasicImageConfig with signature make_image_config(shape, pixel_size, voltage_in_kilovolts).

get_metadata(index: int | slice | Int[ndarray, ''] | Int[ndarray, '_']) -> pandas.DataFrame

Extract the columns at a given index not included in the cryoJAX classes loaded into parameter_info = parameter_file[index].

This method is for advanced usage; prefer setting parameter_file.loads_metadata = True and accessing this as parameter_info = parameter_file[...]; metadata = parameter_info["metadata"].

Datasets: parameter and image manipulation

cryospax.AbstractParticleDataset

cryospax.AbstractParticleDataset(cryospax.AbstractDataset)

Base class for a particle dataset.

parameter_file property

mode property

only_images property

__getitem__(index) -> ~T1

__setitem__(index, value: ~T2)

__len__() -> int

append(value: ~T2)

write_images(index_array: Int[ndarray, '_'], images: Float[NDArrayLike, '... _ _'], parameters: PyTree | None = None)

cryospax.RelionParticleDataset(cryospax.AbstractParticleDataset)

A dataset that wraps a RELION particle stack in STAR format.

__init__(parameter_file: cryospax.AbstractRelionParticleParameterFile, path_to_relion_project: str | pathlib.Path, mode: Literal['r', 'w'] = 'r', *, mrcfile_options: dict[str, Any] = {}, only_images: bool = False)

Arguments:

• path_to_relion_project: In RELION STAR files, only a relative path is added to the 'rlnImageName' column. This is relative to the path to the "project", which is given by this parameter.
• parameter_file: The cryospax.RelionParticleParameterFile.
• mode:
  • If mode = 'w', the dataset is prepared to write new images. This is done by removing 'rlnImageName' from parameter_file.starfile_data, if it exists at all. does not have a column 'rlnImageName' and image files are not yet written.
  • If mode = 'r', images are read from the 'rlnImageName' stored in the parameter_file.starfile_data.
• mrcfile_options: A dictionary with options for MRC file writing. This accepts the following keys:
  • 'prefix': A str which acts as the prefix to the filenames. If this is equal to 'img', then the filename for image stack 0 will be called "img_00000.mrcs", for delimiter = '_' and n_characters = 5. By default, 'img'.
  • 'output_folder': A str or pathlib.Path type where to write MRC files, relative to the path_to_relion_project.
  • 'n_characters': An int for the number of characters to write the filename number string. If this is equal to 5, then the filename for image stack 0 will be called "img_00000.mrcs", for delimiter = '_' and prefix = 'img'. By default, 5.
  • 'delimiter': A str for the delimiter between the filename prefix and number string. If this is equal to '_', then the filename for image stack 0 will be called "img_00000.mrcs", for n_characters = 5 and prefix = 'img'. By default, 'img'.
  • 'overwrite': If True, overwrite existing MRC file path if it exists. By
• only_images: If False, load parameters and images. Otherwise, load only images.

load(path_to_starfile: str | pathlib.Path, path_to_relion_project: str | pathlib.Path, *, selection_filter: dict[str, Callable] = {}, max_optics_groups: int | None = None, mrcfile_options: dict[str, Any] = {}, loads_metadata: bool = False, loads_envelope: bool = False, make_image_config: Callable[[tuple[int, int], FloatLike, FloatLike], cryojax.simulator.BasicImageConfig] = default_fn, only_images: bool = False) -> typing.Self classmethod

Convenience wrapper for loading a cryospax.RelionParticleDataset in mode = 'r' and reading via particle_info = dataset[index] or writing via dataset[index] = particle_info.

See cryospax.RelionParticleParameterFile.__init__ and cryospax.RelionParticleDataset.__init__ for more information.

empty(path_to_starfile: str | pathlib.Path, path_to_relion_project: str | pathlib.Path, num_particles: int, *, max_optics_groups: int = 1, exist_ok: bool = False, mrcfile_options: dict[str, Any] = {}) -> typing.Self classmethod

Convenience wrapper for intializing a new cryospax.RelionParticleDataset in mode = 'w' and writing via dataset[index] = particle_info or dataset.append(particle_info).

See cryospax.RelionParticleParameterFile.__init__ and cryospax.RelionParticleDataset.__init__ for more information.

__getitem__(index: int | slice | Int[ndarray, ''] | Int[ndarray, 'N']) -> dict[str, Any]

Load dataset with value = dataset[...] syntax, where value is a dictionary with keys:

• 'images': An image or image stack to write to an MRC file. This key is required.
• 'parameters': See cryospax.RelionParticleParameterFile for more information. This key is not included if only_images = True.

__setitem__(index: int | slice | Int[ndarray, ''], value: dict[str, Any])

Set dataset entries with dataset[...] = value syntax, where value is a dictionary with keys:

• 'images': An image or image stack to write to an MRC file. This key is required.
• 'parameters': See cryospax.RelionParticleParameterFile for more information. This key is optional.

__len__() -> int

The number of particles in the STAR file.

append(value: dict[str, Any])

Add an entry or entries to the dataset with dataset.append(value) syntax, where value is a dictionary with keys:

• 'images': An image or image stack to write to the MRC file.
• 'parameters': See cryospax.RelionParticleParameterFile for more information.

Both keys are required.

write_images(index_array: Int[ndarray, '_'], images: Float[NDArrayLike, '... _ _'], parameters: dict[str, Any] | None = None)

Write images to the dataset given images and parameters.

This function is wrapped by cryospax.RelionParticleDataset.append and cryospax.RelionParticleDataset.__setitem__ and in most cases these APIs are preferred.

parameter_file property

The parameter_file that reads/writes from the STAR file. This is modified internally during calls to dataset[...] = value.

This cannot be modified after initialization.

path_to_relion_project property

The path to the RELION project. Paths in the STAR file are relative to this directory.

This cannot be modified after initialization.

mode property

Whether or not the dataset was instantiated in reading ('r') or writing ('w') mode.

This cannot be modified after initialization.

particle_data property

The pandas.DataFrame of particle STAR file entries.

Alias to cryospax.AbstractRelionParticleParameterFile.particle_data.

optics_data property

The pandas.DataFrame of optics STAR file entries.

Alias to cryospax.AbstractRelionParticleParameterFile.optics_data.

mrcfile_options property

Settings for writing MRC files with. See cryospax.RelionParticleDataset.__init__ for more information.

only_images property

If True, load images and not parameters. This gives better performance when it is not necessary to load parameters.

dataset.only_images = True
particle_info = dataset[0]
assert "images" in particle_info  # True
assert "parameters" not in particle_info  # True

Basic I/O

cryospax.read_starfile(filename: str | pathlib.Path, **kwargs: Any) -> dict[str, pandas.DataFrame]

Read a STAR file using starfile.

Arguments:

• filename: The path where to read the STAR file. This must include a '.star' extension.

Keyword arguments are passed to starfile.read.

cryospax.write_starfile(starfile_data, filename: str | pathlib.Path, **kwargs: Any)

Write a STAR file using starfile.

Arguments:

• starfile_data: A dictionary whose keys are strings and whose entries are pandas.DataFrames.
• filename: The path where to write the STAR file. This must include a '.star' extension.

Keyword arguments are passed to starfile.write.