Skip to content

recovar.em

Expectation-Maximization algorithms for pose refinement and heterogeneous reconstruction.

states

EM state containers: EMState, SGDState, HeterogeneousEMState.

recovar.em.states

EM state containers: EMState, SGDState, HeterogeneousEMState.

iterations

High-level EM loop orchestration and convergence tracking.

recovar.em.iterations

High-level EM loop orchestration and convergence tracking.

core

Core EM iteration logic: cross-correlation, residual computation.

recovar.em.core

Core EM iteration logic: cross-correlation, residual computation.

compute_dot_products(projections, batch, translations, CTF_params, ctf, noise_variance, process_images, voxel_size, image_shape)

Computes -2 * y_i.T @ (S_s C_i * Proj_j) for i,j,s where C_i is CTF, S_s are shifts, and Proj_j are projections, y_i are batch (unprocessed)

compute_dot_products_eqx(config, projections, batch, translations, ctf_params, noise_variance)

Equinox version of compute_dot_products (9 → 6 params).

compute_CTFed_proj_norms_eqx(config, projections, ctf_params, noise_variance)

Equinox version of compute_CTFed_proj_norms (6 → 4 params).

compute_CTFed_proj_norms(projections, CTF_params, ctf, noise_variance, voxel_size, image_shape)

Computes |C_i Proj_j|^2 for i,j by writing it as a mat-mat where C_i is CTF, S_s are shifts, and Proj_j are projections

e_step

E-step: posterior probability computation over poses and translations.

recovar.em.e_step

E-step: posterior probability computation over poses and translations.

compute_residuals_many_poses_eqx(config, volumes, images, rotation_matrices, translations, ctf_params, noise_variance, translation_fn='fft')

Equinox version of compute_residuals_many_poses (12 → 8 params).

m_step

M-step: volume update via weighted backprojection.

recovar.em.m_step

M-step: volume update via weighted backprojection.

backproject_one_image_eqx(config, probabilities, images_i, rotation_matrices, translations, ctf_params, noise_variance, translation_fn='fft')

Equinox version of backproject_one_image (12 → 8 params).

sum_up_images_fixed_rots_eqx(config, batch, probabilities, translations, rotations, ctf_params, noise_variance, Ft_y=0, Ft_ctf=0)

Equinox version of sum_up_images_fixed_rots (13 → 9 params).