cryolike.likelihoods

cryolike.likelihoods.interface

cryolike.likelihoods.interface.cross_correlation_full_return.compute_cross_correlation_complete(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → CrossCorrelationReturn

cryolike.likelihoods.interface.cross_correlation_full_return.compute_cross_correlation_complete(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[True]) → tuple[cryolike.likelihoods.interface.cross_correlation_full_return.CrossCorrelationReturn, torch.Tensor]

Compute cross-correlation between templates and images, returning a tensor identifying the cross-correlation for each image-template pair at every displacement and rotation, and (optionally) the integrated log likelihood of each image.

Parameters:

comparator (GeneratorType) – A generator which returns one batch of cross-correlation likelihood computation per call
templates (Templates) – A stack of Templates, with the displacements-to-search already configured
images (Images) – A stack of Images
ctf (CTF) – Contrast transfer function, either singleton (applied to every image) or a tensor of per-point values (one grid per image)
precision (Precision) – The precision at which to return results (note that this is separate from the precision used for computing the cross-correlation likelihood)
include_integrated_likelihood (bool) – Whether to include integrated log likelihood in the return

Returns:

The full cross correlation matrix for each image (if no ILL requested) or a tuple of (same, ILL). ILL is represented as a single score per image.

cryolike.likelihoods.interface.cross_correlation_optimal_pose.compute_optimal_pose(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → OptimalPoseReturn

cryolike.likelihoods.interface.cross_correlation_optimal_pose.compute_optimal_pose(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[True]) → tuple[cryolike.likelihoods.interface.cross_correlation_optimal_pose.OptimalPoseReturn, torch.Tensor]

Compute cross-correlation between templates and images, returning a tensor identifying the best cross-correlation for each image-template pair, as well as tensors identifying the identity of the optimal matching template and the x- and y-displacements and rotation producing the optimal match; and (optionally) the integrated log likelihood of each image.

Parameters:

comparator (GeneratorType) – A generator which returns one batch of cross-correlation likelihood computation per call
templates (Templates) – A stack of Templates, with the displacements-to-search already configured
images (Images) – A stack of Images
ctf (CTF) – Contrast transfer function, either singleton (applied to every image) or a tensor of per-point values (one grid per image)
precision (Precision) – The precision at which to return results (note that this is separate from the precision used for computing the cross-correlation likelihood)
include_integrated_likelihood (bool) – Whether to include integrated log likelihood in the return

Returns:

The: score for each image against its optimal template, as well as the optimal template, displacements, and rotation producing this score (if no ILL requested); or a tuple of (same, ILL). ILL is represented as a single score per image.

Return type:

OptimalPoseReturn | tuple(OptimalPoseReturn, torch.Tensor)

cryolike.likelihoods.interface.cross_correlation_optimal_displacement_rotation.compute_optimal_displacement_and_rotation(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → OptimalDisplacementAndRotationReturn

cryolike.likelihoods.interface.cross_correlation_optimal_displacement_rotation.compute_optimal_displacement_and_rotation(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[True]) → tuple[cryolike.likelihoods.interface.cross_correlation_optimal_displacement_rotation.OptimalDisplacementAndRotationReturn, torch.Tensor]

Compute cross-correlation between templates and images, returning a tensor identifying the best cross-correlation for each image-template pair, as well as tensors identifying the the x- and y-displacements and rotation producing the optimal match; and (optionally) the integrated log likelihood of each image.

Parameters:

comparator (GeneratorType) – A generator which returns one batch of cross-correlation likelihood computation per call
templates (Templates) – A stack of Templates, with the displacements-to-search already configured
images (Images) – A stack of Images
ctf (CTF) – Contrast transfer function, either singleton (applied to every image) or a tensor of per-point values (one grid per image)
precision (Precision) – The precision at which to return results (note that this is separate from the precision used for computing the cross-correlation likelihood)
include_integrated_likelihood (bool) – Whether to include integrated log likelihood in the return

Returns:

OptimalDisplacementAndRotationReturn | tuple(OptimalDisplacementAndRotationReturn,: torch.Tensor): The score for each image against each template, as well as the optimal displacements and rotation producing this score (if no ILL requested); or a tuple of (same, ILL). ILL is represented as a single score per image.

cryolike.likelihoods.interface.cross_correlation_optimal_displacement.compute_optimal_displacement(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → OptimalDisplacementReturn

cryolike.likelihoods.interface.cross_correlation_optimal_displacement.compute_optimal_displacement(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[True]) → tuple[cryolike.likelihoods.interface.cross_correlation_optimal_displacement.OptimalDisplacementReturn, torch.Tensor]

Compute cross-correlation between templates and images, returning optimal displacement, and (optionally) the integrated log likelihood for each template-image pair.

Parameters:

comparator (GeneratorType) – A generator which returns one batch of cross-correlation likelihood computation per call
templates (Templates) – A stack of Templates, with the displacements-to-search already configured
images (Images) – A stack of Images
ctf (CTF) – Contrast transfer function, either singleton (applied to every image) or a tensor of per-point values (one grid per image)
precision (Precision) – The precision at which to return results (note that this is separate from the precision used for computing the cross-correlation likelihood)
include_integrated_likelihood (bool) – Whether to include integrated log likelihood in the return

Returns:

The: optimal displacement for each image (if no ILL requested) or a tuple of (same, ILL). ILL is represented as a matrix of scores indexed as [image, template].

Return type:

OptimalDisplacementReturn | tuple(OptimalDisplacementReturn, torch.Tensor)

cryolike.likelihoods.interface.cross_correlation_optimal_rotation.compute_optimal_rotation(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → OptimalRotationReturn

cryolike.likelihoods.interface.cross_correlation_optimal_rotation.compute_optimal_rotation(comparator: GeneratorType, templates: Templates, images: Images, precision: Precision, include_integrated_log_likelihood: Literal[False]) → tuple[cryolike.likelihoods.interface.cross_correlation_optimal_rotation.OptimalRotationReturn, torch.Tensor]

Compute cross-correlation between templates and images, returning score and optimal rotation achieving that score for each image-template pair, and (optionally) the integrated log likelihood for each image.

Parameters:

device (torch.device) – Device to use for computation
images_fourier (torch.Tensor) – Fourier-space images
ctf (torch.Tensor) – Contrast transfer function, either singleton (applied to every image) or a vector (one per image).
n_pixels_phys (int) – Number of pixels in the physical image, i.e. the physical image source’s Cartesian grid
n_images_per_batch (int) – Number of images to use per batch
n_templates_per_batch (int) – Number of templates to use per batch
return_integrated_likelihood (bool) – Whether to include integrated log likelihood in the return

Returns:

The: optimal rotation for each image (if no ILL requested) or a tuple of (same, ILL). ILL is represented as a matrix of scores indexed as [image, template].

Return type:

OptimalRotationReturn | tuple(OptimalRotationReturn, torch.Tensor)

cryolike.likelihoods.iteration

class cryolike.likelihoods.iteration.cross_correlation_iterator_types.CrossCorrelationYieldType(t_start: int, t_end: int, i_start: int, i_end: int, cross_correlation_msdw: Tensor, log_likelihood_ms: Tensor | None)

Bases: NamedTuple

Return type yielded by a cross-correlation generator. These generators are designed to provide an interface between the cross-correlation likelihood computation, on the one hand (which should not have to worry about managing tensor memory or device locations), and collection/summarization on the other hand (which aggregates across different iterations through the large template/image stacks).

t_start

Index (from larger stack) of first template used in this set of comparisons

Type:: int

t_end

1 + index (from larger stack) of last template used in this set of comparisons. Together, the range [t_start:t_end] should be the range, from a larger aggregation tensor, into which this set of comparisons could be copied.

Type:: int

i_start

Index (from larger stack) of first image used in this set of comparisons

Type:: int

i_end

1 + index (from larger stack) of last image used in this set of comparisons (range as with templates)

Type:: int

cross_correlation_msdw

The full results of cross-correlation likelihood, properly normalized by image and template norms, indexed as [image, template, displacement, inplane-rotation]. The m- and s-dimensions should match the sizes suggested by t_start/t_end and i_start/i_end, above.

Type:: torch.Tensor

log_likelihood_ms

If set, the result of computing the integrated log likelihood of this image range against the tensors. Will be None if ill computation was not requested.

Type:: torch.Tensor | None

cross_correlation_msdw: Tensor: Alias for field number 4

i_end: int: Alias for field number 3

i_start: int: Alias for field number 2

log_likelihood_ms: Tensor | None: Alias for field number 5

t_end: int: Alias for field number 1

t_start: int: Alias for field number 0

cryolike.likelihoods.iteration.template_image_comparator.template_first_comparator(device: device, images: Images, templates: Templates, ctf: CTF, n_images_per_batch: int, n_templates_per_batch: int, return_integrated_likelihood: bool, precision: Precision = Precision.DEFAULT) → Generator[CrossCorrelationYieldType, None, None]

Creates a generator that carries out the pairwise image-template comparison to compute cross-correlation likelihood (and optional integrated log likelihood). This version favors templates over images, i.e. it expects the number of templates to be much smaller than the number of images and keeps the templates in memory as much as possible.

Parameters:

device (torch.device) – Device to use to carry out the calculations (usually cpu or cuda)
images (Images) – Stack of images to iterate over
templates (Templates) – Stack of templates to iterate over. We expect that any searched displacements will already have been configured on the Templates object.
ctf (CTF) – Contrast transfer function object to use in adjusting templates
n_images_per_batch (int) – Number of images to use in a single iteration
n_templates_per_batch (int) – Number of templates to use in a single iteration
return_integrated_likelihood (bool) – Whether to include integrated log likelihood
precision (Precision, optional) – Precision to use for the cross-correlation calculation. Defaults to Precision.DEFAULT, in which case we will follow the precision used on the templates.

Yields:

CrossCorrelationYieldType –

For every iteration, we will return (in order) the: start index and end index of the templates used, the start and end indices of the images used, the full cross-correlation likelihood over that range, and the integrated log likelihood of the images (if requested, else None). These are strongly typed in a NamedTuple for ease of use.

cryolike.likelihoods.iteration.integrated_log_likelihood_aggregation.aggregate_ill(templates: Templates, log_likelihood_MS: Tensor, precision: Precision) → Tensor

Aggregates the per-image-template-pair integrated log likelihood into a single per-image likelihood using logsumexp.

Parameters:

templates (Templates) – Template stack from the comparison, used to find the viewing-angle weights and the number of displacements (which assumes the templates are displaced; if the images are displaced, we will need to add an additional parameter)
log_likelihood_MS (torch.Tensor) – The per-pair log likelihood values
precision (Precision) – Precision for the computation

Returns:

An aggregate log likelihood per image

Return type:

torch.Tensor

cryolike.likelihoods.kernels

cryolike.likelihoods.kernels.cross_correlation_likelihood_kernel.compute_cross_correlation(n_inplanes: int, sqrtweighted_premultiplied_CTF_image_fourier_bessel_conj_mnq: Tensor, sqrtweighted_displaced_fourier_bessel_templates_sdnq: Tensor) → Tensor

Computes the raw cross-correlation between image and template, by doing a dot product on the fourier-bessel-space representations of the image and template stacks. We expect that the image stack will be a 3-tensor indexed as [image, grid-point, frequency-dimension] and the template stack will be a 4-tensor indexed as [template, displacement, grid-point, frequency].

Note that both image and template stacks are expected to have been multiplied by the square root of the integration weights for the polar grid, and that the images should have been further multiplied by the CTF.

Parameters:

n_inplanes (int) – Number of inplane rotations in the grid (needed to invert the bessel-space transform)
sqrtweighted_premultiplied_CTF_image_fourier_bessel_conj_mnq (torch.Tensor) – A fourier-bessel representation of the image stack, with its points pre-multiplied by the square root of the quadrature weights and by the CTF
sqrtweighted_displaced_fourier_bessel_templates_sdnq (torch.Tensor) – A fourier-bessel representation of the templates, premultiplied by the square root of the quadrature weights, and displaced across the searched displacement grid (as a separate dimension)

Returns:

The unnormalized cross-correlation between images and templates,: indexed as [image, template, displacement, inplane-rotation]

Return type:

torch.Tensor

cryolike.likelihoods.kernels.cross_correlation_likelihood_kernel.compute_image_norms(images_fourier: Tensor)

Computes the norms of each image. This is currently fairly straightforward since they do not require displacement, rotation, or CTF-application.

Parameters:

images_fourier (torch.Tensor) – Tensor of the image stack (indexed as [image, radius, inplane-rotation] i.e. on the polar quadrature grid). The image values should have been premultiplied by the square root of the quadrature grid weights.

Returns:

The per-image norm, with the tensor expanded to include: size-1 indices for the template, displacement, and rotation dimensions

Return type:

torch.Tensor

cryolike.likelihoods.kernels.cross_correlation_likelihood_kernel.compute_template_norms(n_inplanes: int, sqrtweighted_fourier_templates_snw: Tensor, ctf_mnw: Tensor)

Computes the (CTF-adjusted) norms for each template at every inplane rotation (which matters for anisotropic CTFs). This is done in Fourier-Bessel representation, which avoids having to materialize all the different template rotations.

Parameters:

n_inplanes (int) – Number of inplane rotations in the polar representation (needed to invert the fourier-bessel transform)
sqrtweighted_fourier_templates_snw (torch.Tensor) – A tensor representing the template stack, indexed as [template, grid-point, inplane-rotation]. Note that this does not need to be displaced as template norms should not be affected by displacement. The point values should have been premultiplied by the square root of the quadrature point weights.
ctf_mnw (torch.Tensor) – A tensor of the CTF values for the images, indexed as [image/ctf id, linear grid point, inplane-rotation]

Returns:

A tensor of the template norms, indexed as: [image, template, inplane-rotation] (where the image index is 1:1 with the CTF index)

Return type:

torch.Tensor

cryolike.likelihoods.kernels.integrated_log_likelihood_kernel.ill_kernel(Iss: float, n_pixels_phys: int, sqrt_weighted_mask_points: Tensor, CTF_sqrtweighted_fourier_templates_msnw: Tensor, images_fourier_sqrtweighted_mnw: Tensor, Ixx: Tensor, Iyy: Tensor, Ixy: Tensor)

Compute the per-image integrated log likelihood.

Parameters:

Iss (float) – Scalar value resulting from integrating the mask weights
n_pixels_phys (int) – Total number of pixels in the real-space representation of the images
sqrt_weighted_mask_points (torch.Tensor) – Mask points from the quadrature grid, pre-multiplied by the square root of the quadrature weights
CTF_sqrtweighted_fourier_templates_msnw (torch.Tensor) – Template stack in Fourier-space representation, pre-multiplied by the per-image CTF and the square root of the quadrature weights, and indexed as [image/ctf, template, radius, inplane-angle]
images_fourier_sqrtweighted_mnw (torch.Tensor) – Image stack in Fourier-space representation, pre-multiplied by the square root of the quadrature weights, and indexed as [image, radius, inplane-angle]
Ixx (torch.Tensor) – Template norms, as a 4-tensor [image, template, displacement, inplane]
Iyy (torch.Tensor) – Image norms, as a 4-tensor [image, template, displacement, inplane]
Ixy (torch.Tensor) – Unnormalized cross-correlation dot product between templates and images, indexed as [image, template, displacement, inplane-rotation]

Returns:

The integrated log likelihood of each image-tensor pair, indexed: as [image, tensor]

Return type:

torch.Tensor