tomodrgn.commands.downsample.collate_particle_tilts

collate_particle_tilts(batch: list[tuple[ndarray, ndarray]]) → tuple[tensor, tensor]

Collate a batch of particle images, optionally with a variable number of tilt images associated with each particle. Given input of shape ([ind1], [a, X, X]), ([ind2], [b, X, X]) where a and b can be any number, returns [ind1, ind2], [[a+b, X, X]].

Parameters:

batch – list of 2-tuples of particle indices (np.ndarray of shape (ntilts)) and associated images (np.ndarray of shape (ntilts, npixels, npixels)

Returns:

2-tuple of concatenated particle indices and concatenated particle images