tomodrgn.models.VolumeGenerator#

class VolumeGenerator(config: str | dict[str, dict], weights_path: str | None = None, model: TiltSeriesHetOnlyVAE | FTPositionalDecoder | None = None, lat: Lattice | None = None, amp: bool = True)[source]#

Bases: object

Convenience class to generate volume(s) from a trained tomoDRGN model. Supports evaluating homogeneous (train_nn) or heterogeneous (train_vae) models.

Methods

generate_volumes

Generate volumes at specified latent embeddings and save to specified output directory.
generate_volumes(z: ndarray | str | None, out_dir: str, out_name: str = 'vol', downsample: int | None = None, lowpass: float | None = None, flip: bool = False, invert: bool = False, batch_size: int = 1) → None[source]#

Generate volumes at specified latent embeddings and save to specified output directory. Generated volume filename format is out_dir / out_name _ {i:03d if multiple volumes} .mrc

Parameters:

  • z – latent embeddings to evaluate as numpy array of shape (nptcls, zdim), or path to a .txt or .pkl file containing array of shape (nptcls, zdim). None if evaluating a homogeneous tomodrgn model.

  • out_dir – path to output directory in which to save output mrc file(s)

  • out_name – string to prepend to output .mrc file name(s)

  • downsample – downsample reconstructed volumes to this box size (units: px) by Fourier cropping, None means to skip downsampling

  • lowpass – lowpass filter reconstructed volumes to this resolution (units: Å), None means to skip lowpass filtering

  • flip – flip the chirality of the reconstructed volumes by inverting along the z axis

  • invert – invert the data sign of the reconstructed volumes (light-on-dark vs dark-on-light)

  • batch_size – batch size to parallelize volume generation (32-64 works well for box64 volumes)

Returns:

None