tomodrgn.commands.train_nn.decode_batch#

decode_batch(*, model: FTPositionalDecoder | DataParallelPassthrough, lat: Lattice, batch_rots: Tensor, batch_hartley_2d_mask: Tensor) → Tensor[source]#

Decode a batch of particles represented by multiple images from per-particle latent embeddings and corresponding lattice positions to evaluate

Parameters:

• model – FTPositionalDecoder object to be trained

• lat – Hartley-transform lattice of points for voxel grid operations

• batch_rots – Batch of 3-D rotation matrices corresponding to batch_images known poses, shape (batchsize, ntilts, 3, 3)

• batch_hartley_2d_mask – Batch of 2-D masks to be applied per-spatial-frequency, shape (batchsize, ntilts, boxsize_ht**2) Calculated as the intersection of critical dose exposure curves and a Nyquist-limited circular mask in reciprocal space, including masking the DC component.

Returns:

Reconstructed central slices of Fourier space volumes corresponding to each particle in the batch, shape (batchsize * ntilts * boxsize_ht**2 [batch_hartley_2d_mask]). Note that the returned array is completely flattened (including along batch dimension) due to potential for uneven/ragged reconstructed image tensors per-particle after masking