recovar.core¶

Low-level primitives for cryo-EM forward modeling, geometry, CTF handling, and Fourier-space operations.

configs¶

Equinox-based configuration containers for the forward model, batch data, and model state.

`recovar.core.configs` ¶

Equinox-based configuration modules for clean parameter passing.

These modules bundle the many parameters that are passed to JAX-jitted functions throughout RECOVAR into structured, typed containers. This eliminates the need for complex static_argnums lists and keeps function signatures readable.

`ForwardModelConfig` ¶

Bases: Module

Bundles geometry, CTF, and discretization parameters.

All fields are static (compile-time constants), so changing any value triggers JAX recompilation — but these values are fixed for a given dataset and reconstruction run.

`base_volume_shape` `property` ¶

Original (non-upsampled) volume shape.

`compute_ctf(ctf_params, *, half_image=False)` ¶

Compute CTF values for a batch of images.

Parameters¶

half_image : if True, evaluate on the rfft-packed half-spectrum grid.

`compute_ctf_half(ctf_params)` ¶

Convenience alias for compute_ctf(ctf_params, half_image=True).

`compute_ctf_at_shape(ctf_params, image_shape)` ¶

Compute CTF on a different frequency grid (e.g. upsampled).

`replace(**kwargs)` ¶

Create a new config with some fields replaced.

Useful for changing static fields (e.g. disc_type) which cannot be modified via eqx.tree_at.

`from_dataset(cryo, disc_type='linear_interp', process_fn=None, upsampling_factor=None)` `classmethod` ¶

Create from a CryoEMDataset instance.

Parameters¶

cryo : CryoEMDataset Source dataset for geometry and CTF configuration. disc_type : str Discretization type (e.g. 'linear_interp', 'cubic', ''). process_fn : callable, optional Image preprocessing function applied inside jitted code. upsampling_factor : int, optional Volume upsampling factor (e.g. 2 for 2× oversampled grid). Computes the upsampled volume shape directly without mutating the dataset object.

`ModelState` ¶

Bases: Module

Current reconstruction state passed to jitted functions.

Contains the mean estimate, volume mask, and optionally the PCA basis and eigenvalues used for covariance/embedding computations.

`CovColumnOpts` ¶

Bases: Module

Static options for covariance column (H/B) computation.

`CovarianceOpts` ¶

Bases: Module

Options for covariance estimation inner functions.

`EmbeddingOpts` ¶

Bases: Module

Options for embedding / latent coordinate computation.

ctf¶

CTF evaluation and parameter handling.

`recovar.core.ctf` ¶

CTF evaluation and parameter handling for cryo-EM images.

Provides :class:CTFEvaluator, a unified equinox module dispatched by :class:CTFMode that evaluates the Contrast Transfer Function for all supported cryo-EM/ET imaging modes.

`CTFParamIndex` ¶

Bases: IntEnum

Enum for CTF parameter indices to avoid magic numbers.

`CTFMode` ¶

Bases: IntEnum

CTF evaluation mode, determined by the imaging modality.

Attributes:

Name	Type	Description
`SPA`		Standard single-particle analysis CTF.
`SPA_ANTIALIASED`		SPA with 2x-upsampled antialiasing filter.
`TILT_SERIES`		Parametric dose weighting from `dose_per_tilt` and `angle_per_tilt` constants.
`CRYO_ET`		Per-image dose and tilt-angle columns already present in the CTF parameter array.

`CTFEvaluator` ¶

Bases: Module

Unified CTF evaluator for all cryo-EM/ET imaging modes.

An equinox module that is callable with the standard CTF signature::

evaluator(ctf_params, image_shape, voxel_size, *, half_image=False)

All fields are static (compile-time constants). Changing mode or dose parameters triggers JAX recompilation, which is the correct behaviour since these are fixed per dataset.

Parameters¶

mode : CTFMode Imaging modality. Defaults to :attr:CTFMode.SPA. dose_per_tilt : float Dose per tilt in e-/A^2. Only used when mode == TILT_SERIES. angle_per_tilt : float Tilt-angle increment in degrees. Only used when mode == TILT_SERIES.

`call(ctf_params, image_shape, voxel_size, *, half_image=False)` ¶

Evaluate the CTF for a batch of images.

Parameters¶

ctf_params : jax.Array Per-image CTF parameters, shape (N, K). image_shape : tuple[int, int] Image dimensions in pixels. voxel_size : float Pixel size in Angstroms. half_image : bool If True, evaluate on the rfft-packed half-spectrum grid.

`as_ctf_evaluator(fn_or_evaluator)` ¶

Coerce a callable into a CTFEvaluator-compatible eqx.Module.

If fn_or_evaluator is already a :class:CTFEvaluator or :class:_LegacyCTFAdapter, return it unchanged. Otherwise wrap it in a :class:_LegacyCTFAdapter.

`evaluate_ctf(freqs, ctf_params)` ¶

Evaluate the Contrast Transfer Function for a batch of images.

Broadcasts per-image parameters over a shared frequency grid in a single fused kernel (no vmap).

Parameters¶

freqs : array (n_pixels, 2) Shared frequency coordinates in 1/Angstrom. ctf_params : array (n_images, K) Packed CTF parameters per image. Layout: [DFU, DFV, DFANG, VOLT, CS, W, PHASE_SHIFT, BFACTOR, CONTRAST, ...] — see :class:CTFParamIndex.

Returns¶

array (n_images, n_pixels)

`compute_antialiased_ctf_squared(ctf_fn, ctf_params, image_shape, voxel_size, *, half_image=False)` ¶

Compute CTF² with 2x-upsampled antialiasing, in full or half-image format.

Evaluates CTF on a 2x grid, squares, box-filters, and downsamples. This reduces aliasing in the Wiener weight accumulator.

geometry¶

Image translation, rotation, and coordinate transforms.

`recovar.core.geometry` ¶

Image translation, rotation matrix construction, and coordinate transforms.

`get_gridpoint_coords(rotation_matrix, image_shape, volume_shape)` ¶

Compute 3-D Fourier grid coordinates for a rotated central slice.

Applies rotation_matrix to the unrotated 2-D frequency plane and returns the corresponding 3-D grid indices in the volume.

Parameters:

Name	Description	Default
`rotation_matrix`	Rotation matrix, shape `(3, 3)`.	required
`image_shape`	2-D image side length (static).	required
`volume_shape`	3-D volume shape tuple (static).	required

Returns:

Type	Description
	Rotated grid coordinates, shape `(image_size, 3)`.

`get_unrotated_half_plane_grid_points(image_shape, three_d_upsampling_factor=1)` ¶

Unrotated 2-D frequency grid for only the H*(W//2+1) half-image pixels.

Like :func:get_unrotated_plane_grid_points but generates coordinates directly for only the non-redundant rfft-packed pixels.

Pixel ordering matches :func:~recovar.core.fourier_transform_utils.full_image_to_half_image: uses _half_image_pixel_indices to select the same packed indices as full_image_to_half_image. This keeps behavior correct for non-square frequency grids.

`translate_images(image, translation, image_shape, *, half_image=False)` ¶

Apply in-plane translations to Fourier-space images via phase shifts.

Parameters:

Name	Description	Default
`image`	Fourier-space image(s), shape `(batch, image_size)` or `(image_size,)`.	required
`translation`	Shift(s) in fractional pixel units, shape `(batch, 2)` or `(2,)`.	required
`image_shape`	Image side length (static).	required
`half_image`	If True, use rfft-packed half-spectrum frequency grid.	`False`

Returns:

Type	Description
	Translated image(s) with the same shape as image.

forward¶

Forward model and adjoint operations for cryo-EM image formation.

`recovar.core.forward` ¶

Forward model and adjoint operations for cryo-EM image formation.

Provides Equinox-based APIs (forward_model, adjoint_forward_model, etc.) that take a :class:~recovar.configs.ForwardModelConfig as first argument.

`forward_model(config, volume, ctf_params, rotation_matrices, skip_ctf=False, half_image=False, half_volume=False)` ¶

Project volume into images via slice-and-CTF forward model.

Parameters¶

half_image : bool If True, return rfft-packed half-spectrum images and use config.compute_ctf_half for CTF, roughly halving memory and compute. half_volume : bool If True, volume is an rfft-packed half-volume (N0*N1*(N2//2+1),).

`forward_model_and_adjoint(config, volume, ctf_params, rotation_matrices, skip_ctf=False)` ¶

Forward model plus its VJP (adjoint) closure.

`adjoint_forward_model(config, slices, ctf_params, rotation_matrices, skip_ctf=False, volume=None, half_image=False, half_volume=False)` ¶

Adjoint of the forward model (direct back-projection).

Uses :func:adjoint_slice_volume which dispatches to CUDA when available, avoiding the VJP-through-FFI overhead.

Parameters¶

volume : optional accumulator array to add into (avoids fresh allocation). half_image : if True, slices are rfft-packed half-spectrum images. CTF is computed in half-spectrum format when skip_ctf=False. half_volume : if True, output volume uses rfft-packed layout (N0 * N1 * (N2 // 2 + 1),). Only supported with CUDA.

`compute_AtAv(config, volume, ctf_params, rotation_matrices, noise_variance, skip_ctf=False)` ¶

Compute A^T (A v / noise_variance) for normal equations.

slicing¶

Fourier slice extraction and backprojection.

`recovar.core.slicing` ¶

Fourier slice extraction and adjoint (backprojection) operations.

Dispatch rules

GPU + order 0/1 → CUDA project + CUDA backproject
GPU + order 3 → CUDA project + JAX VJP backproject (cubic)
CPU + order 0/1 → RELION-style JAX (explicit scatter backproject)
CPU + order 3 → JAX map_coordinates + VJP backproject

Three core public functions handle all volume/image format combinations via half_volume and half_image parameters: - :func:slice_volume (forward projection) - :func:batch_slice_volume (batched forward) - :func:adjoint_slice_volume (backprojection)

CUDA custom-VJP wrappers live in :mod:recovar.core.cuda_ops. RELION-style JAX reference: :mod:recovar.core.relion_interp.

`slice_volume(volume, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO)` ¶

Project volume to images via interpolation.

Parameters¶

volume : array, flat (vol_size,) or grid-shaped. Real or complex. Promoted to complex for the CUDA kernel. half_volume : if True, volume is rfft-packed (N0*N1*(N2//2+1),). half_image : if True, output images are rfft-packed (n, H*(W//2+1)). max_r : sphere clipping radius. Default (_AUTO) uses image_shape[0]//2 - 1. Pass None to disable clipping.

`batch_slice_volume(volumes, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO)` ¶

Project a batch of volumes to images.

Parameters¶

volumes : array (batch, vol_size). Real or complex. Promoted to complex for the CUDA kernel. half_volume : if True, volumes are rfft-packed half-volumes. half_image : if True, output images are rfft-packed. max_r : sphere clipping radius. Default uses image_shape[0]//2 - 1. Pass None to disable.

`adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volume=None, half_image=False, half_volume=False, max_r=_AUTO)` ¶

Adjoint slice extraction (backprojection).

Parameters¶

slices : array (n_images, n_pixels). Real or complex. If volume is complex, slices are promoted to match (and vice versa). half_image : if True, slices are rfft-packed half-spectrum images. half_volume : if True, output uses rfft-packed half-volume layout. volume : optional accumulator to add the result into. max_r : sphere clipping radius. Default uses image_shape[0]//2 - 1. Pass None to disable.

`batch_adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volumes=None, half_image=False, half_volume=False, max_r=_AUTO)` ¶

Batch backprojection: per-volume image sets to batch of volumes.

Parameters¶

slices : shape (batch, n_images, n_pixels). Real or complex. Promoted to match volumes dtype (and vice versa). rotation_matrices : shape (n_images, 3, 3) — shared across batch. volumes : optional (batch, vol_flat_size) accumulators. half_image, half_volume : same semantics as adjoint_slice_volume. max_r : sphere clipping radius. Default uses image_shape[0]//2 - 1. Pass None to disable.

`precompute_cubic_coefficients(volume, volume_shape)` ¶

Precompute periodic cubic B-spline coefficients for a full volume.

Uses periodic boundary conditions (FFT-based circulant solve). Output has the same shape as input (no boundary padding).

`precompute_cubic_coefficients_half(volume, volume_shape)` ¶

Precompute periodic cubic coefficients, stored as half-volume.

Since periodic coefficients preserve Hermitian symmetry, this is lossless.

`slice_from_cubic_coefficients(coeffs, rotation_matrices, image_shape, volume_shape, half_image=False)` ¶

Project from precomputed periodic cubic coefficients to images.

Parameters¶

coeffs : precomputed spline coefficients from :func:precompute_cubic_coefficients. Can be full-volume (N0,N1,N2) or half-volume (N0N1(N2//2+1),) layout. half_image : if True, output images are rfft-packed (n, H*(W//2+1)).

indexing¶

Frequency-to-volume index mapping and radial shell utilities.

`recovar.core.indexing` ¶

Frequency-to-volume index mapping and radial shell utilities.

mask¶

Real-space and Fourier-space mask generation and manipulation.

`recovar.core.mask` ¶

Real-space and Fourier-space mask generation and manipulation.

Provides masks for the reconstruction pipeline:

Volume masks: make_mask (standalone, RELION-style), make_mask_from_half_maps (averages half-maps then calls make_mask), masking_options (pipeline entry point), make_mask_from_gt, make_union_gt_mask
Radial / spherical masks: get_radial_mask, raised_cosine_mask, soft_mask_outside_map
Image masks: window_mask, smooth_circular_mask
Softening: soften_volume_mask (distance-transform cosine taper)

`masking_options(volume_mask_option, means, volume_shape, dtype_real=np.float32, mask_dilation_iter=0, keep_input_mask=False, dilated_mask_dilations_iter=None)` ¶

Build volume mask and dilated mask from a pipeline mask specification.

Parameters:

Name	Description	Default
`volume_mask_option`	One of: path to `.mrc`, `"from_halfmaps"`, `"sphere"`, `"none"`.	required
`means`	Object with `.corrected0reg` / `.corrected1reg` attributes (used for `from_halfmaps`).	required
`volume_shape`	3-tuple giving the grid dimensions.	required
`dtype_real`	Output dtype.	`float32`
`mask_dilation_iter`	Extra dilation iterations for the input mask.	`0`
`keep_input_mask`	If True, use the MRC mask as-is (no softening).	`False`
`dilated_mask_dilations_iter`	Dilation iterations for the dilated mask. Defaults to `ceil(6 * N / 128)` where N is the grid size.	`None`

Returns:

Type	Description
	`(volume_mask, dilated_volume_mask)` as float arrays in [0, 1].

`make_mask(volume, *, threshold='auto', lowpass_sigma=None, extend=None, soft_edge=3, cleanup=True)` ¶

Create a solvent mask from a 3-D real-space volume.

Follows the same conceptual pipeline as RELION's relion_mask_create: low-pass filter → threshold → (optional cleanup) → extend → soft cosine edge.

Key differences from RELION:

Gaussian low-pass in real space (RELION uses a Fourier-space raised-cosine).
"auto" threshold via Otsu's method (RELION requires a user-specified density threshold, default 0.02 in postprocessing).
Optional morphological cleanup (fill holes, keep largest connected component) — RELION does not do this.
Extension uses distance_transform_edt for spherical (Euclidean) dilation, matching RELION's brute-force Euclidean approach.
Soft edge uses distance_transform_edt + cosine (RELION computes min-Euclidean-distance per voxel — mathematically equivalent).

This function is the building block for :func:make_mask_from_half_maps and can also be called directly on any volume (e.g. a consensus reconstruction loaded from MRC).

Parameters:

Name	Description	Default
`volume`	3-D real-space array (e.g. averaged half-maps, or a single reconstruction).	required
`threshold`	How to binarize. `"auto"` (default): Otsu's method on voxels inside the radial mask. A `float`: fixed density threshold (like RELION's `--ini_threshold`, default 0.02 in postprocessing).	`'auto'`
`lowpass_sigma`	Gaussian sigma in voxels for low-pass smoothing before thresholding. `None` = auto (`max(2, N // 64)`). Set to `0` to disable.	`None`
`extend`	Extension (dilation) of the binary mask in voxels. `None` = auto (`ceil(6 * N / 128)`). (RELION postprocessing default: 3 pixels.)	`None`
`soft_edge`	Width of the cosine soft edge in voxels. (RELION postprocessing default: 6 pixels.)	`3`
`cleanup`	If `True`, fill holes and keep only the largest connected component after thresholding (RELION does not do this).	`True`

Returns:

Type	Description
	Soft mask as a float32 array in [0, 1].

Example::

import mrcfile
from recovar.core.mask import make_mask

h1 = mrcfile.open("half1.mrc").data
h2 = mrcfile.open("half2.mrc").data
avg = (h1 + h2) / 2.0

# Automatic (recommended)
mask = make_mask(avg)

# RELION-like with manual parameters
mask = make_mask(avg, threshold=0.02, extend=3, soft_edge=6,
                 lowpass_sigma=3, cleanup=False)

`make_mask_from_half_maps(halfmap1, halfmap2, smax=3, method='auto', **kwargs)` ¶

Generate a solvent mask from two real-space half-map volumes.

Averages the two half-maps and delegates to :func:make_mask. All keyword arguments are forwarded to make_mask.

For the legacy EMDA local-correlation method, pass method="local_correlation".

Parameters:

Name	Description	Default
`halfmap1, halfmap2`	Real-space half-map volumes (same shape).	required
`smax`	Kernel radius in pixels (only for `method="local_correlation"`).	`3`
`method`	`"auto"` or `"local_correlation"`.	`'auto'`
`**kwargs`	Forwarded to :func:`make_mask` (`threshold`, `lowpass_sigma`, `extend`, `soft_edge`, `cleanup`).	`{}`

Returns:

Type	Description
	Soft mask as a float array in [0, 1].

`make_mask_from_gt(gt_map, smax=3, iter=10, from_ft=True)` ¶

Generate a mask from a ground-truth volume.

Parameters:

Name	Description	Default
`gt_map`	Ground-truth volume (Fourier or real space).	required
`smax`	Kernel radius for thresholding.	`3`
`iter`	Dilation iterations.	`10`
`from_ft`	If True, `gt_map` is in Fourier space.	`True`

Returns:

Type	Description
	Soft mask as a float array in [0, 1].

`make_union_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3)` ¶

Create a union mask from multiple ground-truth real-space volumes.

For each volume, generates a per-volume mask via make_mask_from_gt, thresholds at 0.5, then takes the logical OR of all per-volume masks. The union is dilated and softened to produce the final mask.

Parameters:

Name	Description	Default
`gt_volumes_real`	Either a list of 3-D arrays or a 2-D array of shape `(n_vols, n_voxels)` (reshaped internally to 3-D).	required
`volume_shape`	Tuple giving the 3-D grid dimensions.	required
`smax`	Gaussian kernel radius for `make_mask_from_gt`.	`3`
`iter`	Dilation iterations inside `make_mask_from_gt`.	`1`
`dilation_iters`	Additional dilation iterations applied to the union mask. Defaults to `ceil(6 * volume_shape[0] / 128)` (pipeline convention).	`None`
`kern_rad`	Kernel radius for `soften_volume_mask`.	`3`

Returns:

Type	Description
	Tuple `(soft_mask, binary_mask)` where soft_mask is a float array
	in [0, 1] and binary_mask is the pre-softening boolean array.

`soften_volume_mask(binary_volume_mask, kern_rad=3)` ¶

Soften a binary mask with a cosine taper based on distance transform.

Voxels inside the mask get value 1; voxels within kern_rad of the boundary get a raised-cosine transition; voxels further out get 0.

Parameters:

Name	Type	Description	Default
`binary_volume_mask`		Binary mask (values near 0 or 1).	required
`kern_rad`		Width of the cosine transition in voxels.	`3`

Returns:

Type	Description
	Soft mask as float32 array in [0, 1].

`get_radial_mask(shape, radius=None)` ¶

Binary spherical mask.

Parameters:

Name	Type	Description	Default
`shape`		Volume or image shape tuple.	required
`radius`		Mask radius in voxels. Defaults to `shape[0] // 2 - 1`.	`None`

Returns:

Type	Description
	Boolean array with True inside the sphere.

`raised_cosine_mask(volume_shape, radius, radius_p, offset)` ¶

3D raised-cosine mask with adjustable center.

Value is 1 for r < radius, cosine-tapered for radius <= r < radius_p, and 0 beyond.

Parameters:

Name	Description	Default
`volume_shape`	3-tuple of grid dimensions.	required
`radius`	Inner radius (full value).	required
`radius_p`	Outer radius (zero value).	required
`offset`	3-element center offset.	required

Returns:

Type	Description
	Mask array of shape `volume_shape`.

`soft_mask_outside_map(vol, radius=-1, cosine_width=3, Mnoise=None)` ¶

Soft mask outside map, adapted from RELION.

Applies a raised-cosine mask to vol, replacing voxels outside the mask with the mean background value (or Mnoise if provided).

Parameters:

Name	Description	Default
`vol`	Input volume (JAX or numpy array).	required
`radius`	Mask radius in voxels. Defaults to half the box size.	`-1`
`cosine_width`	Width of the cosine transition band.	`3`
`Mnoise`	Optional replacement value for masked-out regions.	`None`

Returns:

Type	Description
	`(masked_vol, mask)` tuple.

`window_mask(D, in_rad, out_rad)` ¶

2D circular window mask with linear taper (normalised coordinates).

Parameters:

Name	Description	Default
`D`	Image size in pixels (must be even).	required
`in_rad`	Inner radius in normalised coordinates (1.0 = edge).	required
`out_rad`	Outer radius in normalised coordinates.	required

Returns:

Type	Description
	Float32 mask of shape `(D, D)` with values in [0, 1].

`smooth_circular_mask(image_size, radius, thickness)` ¶

2D circular mask with a raised-cosine transition band.

Values are 1 inside radius, 0 outside radius + thickness, and follow a cosine taper in between.

Parameters:

Name	Description	Default
`image_size`	Image size in pixels.	required
`radius`	Inner radius in pixels.	required
`thickness`	Width of the cosine transition in pixels.	required

Returns:

Type	Description
	Float mask of shape `(image_size, image_size)`.

`make_soft_edged_kernel(r1, shape)` ¶

Create a soft-edged convolution kernel for local correlation.

Adapted from EMDA (https://gitlab.com/ccpem/emda).

Parameters:

Name	Type	Description	Default
`r1`		Kernel radius in pixels.	required
`shape`		Volume shape for coordinate generation.	required

Returns:

Type	Description
	Normalised soft-edged kernel.

`threshold_map(arr, prob=0.99, dthresh=None)` ¶

Zero out voxels below the prob percentile threshold.

`make_moving_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3)` ¶

Create a mask for the moving region across GT volumes.

linalg¶

Linear algebra helpers (batch SVD, QR, eigendecomposition).

`recovar.core.linalg` ¶

Linear algebra helpers: batch SVD, QR, eigendecomposition on CPU/GPU.

`thin_svd_in_blocks(X, np=np, memory_to_use=5, epsilon=1e-08, n_components=-1)` ¶

This is an unstable method for SVD but that can run on GPU. This should be only (maybe) used for matrices that cannot fit on the GPU.

Same as SVD but dispatch matrices in blocks to gpu (also pretty unstable)

`thin_svd(X, np=np, epsilon=1e-08)` ¶

For some reason, the built in svd seems to allocate a lot more memory than necessary.

`randomized_svd(A, n_pcs=200)` ¶

Randomized SVD for large matrices that don't fit on GPU.

`broadcast_dot(x, y)` ¶

Batched conjugate inner product: sum_k conj(x[...,k]) * y[...,k].

`broadcast_outer(x, y)` ¶

Batched outer product: x[..., i] * conj(y[..., j]).

`inner_product(x, y)` ¶

Conjugate inner product for non-batched arrays of identical shape.

`batch_inner_product(x, y)` ¶

Conjugate inner product over batch dimension 0.

Accepts arrays shaped (B, ...) and returns (B,).

`half_spectrum_last_axis_weights(last_axis_size, dtype=jnp.float32)` ¶

Weights to recover full-spectrum inner products from packed real FFT coefficients.

`rfft2_hermitian_weights(image_shape, dtype=jnp.float32)` ¶

Precompute sqrt(w) weights for 2-D half-spectrum (rfft2) inner products.

For Hermitian-symmetric arrays A, B of shape (H, W) stored in rfft2 half-spectrum format (shape (H, W//2+1)), the weighted inner product over the half equals the full-spectrum inner product::

sum_{k in half} w[k] * conj(A[k]) * B[k]  ==  sum_{k in full} conj(A[k]) * B[k]

Built from :func:half_spectrum_last_axis_weights, then tiled over the H rows.

Parameters:

Name	Type	Description	Default
`image_shape`		`(H, W)` tuple.	required
`dtype`		dtype for the returned array (default float32).	`float32`

Returns:

Type	Description
	JAX array of shape `(H * (W // 2 + 1),)` with values in `{1, sqrt(2)}`.

`half_spectrum_inner_product(x_half, y_half, full_shape)` ¶

Inner product from packed half-spectrum inputs, equivalent to full-spectrum inner product.

This computes the weighted packed-spectrum sum directly (no full-spectrum reconstruction), applying factor-2 to non-edge frequencies on the packed axis.

`batch_half_spectrum_inner_product(x_half, y_half, full_shape)` ¶

Batched inner products from packed half-spectrum inputs.

Equivalent to computing full-spectrum batched inner products.

fourier_transform_utils¶

FFT utilities and frequency grid construction.

`recovar.core.fourier_transform_utils` ¶

`get_1d_frequency_grid_rfft(n, voxel_size=1, scaled=False)` ¶

Frequency grid for Hermitian-packed real FFT axis.

Returns non-negative bins [0, 1, ..., n//2] (or scaled equivalent).

`get_k_coordinate_of_each_pixel_real(image_shape, voxel_size, scaled=True)` ¶

2D packed-spectrum frequency coordinates for real-input FFT.

image_shape is the original real image shape (H, W). The returned coordinates enumerate the packed spectrum with shape (H, W//2 + 1).

`get_k_coordinate_of_each_pixel_3d_real(image_shape, voxel_size, scaled=True)` ¶

3D packed-spectrum frequency coordinates for real-input FFT.

image_shape is the original real volume shape (D1, D2, D3). The returned coordinates enumerate the packed spectrum with shape (D1, D2, D3//2 + 1).

`get_grid_of_radial_distances_real(image_shape, voxel_size=1, scaled=False, frequency_shift=0, rounded=True)` ¶

Radial distance grid for Hermitian-packed real FFT spectra.

image_shape is the original real-space shape; output shape is the packed spectrum shape with last axis N//2 + 1.

`get_real_fft_packed_shape(shape)` ¶

Return Hermitian-packed spectrum shape for a real-valued signal.

The last axis is reduced from N to N//2 + 1 as in rfft/rfftn.

`get_real_fft_packed_last_axis_indices(n)` ¶

Indices in shifted full spectrum that correspond to packed real bins.

`get_k_coordinate_of_each_pixel_half(image_shape, voxel_size, scaled=True)` ¶

Half-image frequency coords consistent with full_image_to_half_image.

`get_shifted_conjugate_partner_indices(n)` ¶

Index of the conjugate-symmetric partner for each bin in a shifted FFT axis.

For a length-n axis after fftshift, bin i holds frequency u = (i + n//2) % n. Its conjugate partner is at frequency -u % n, which maps back to shifted index (-u % n - n//2) % n. This function returns that mapping as an int32 array of length n.

`get_real_fft_memory_saving_ratio(shape)` ¶

Ratio of packed (rfft) to full spectrum size: prod(packed) / prod(shape).

`image_shape_to_half_image_shape(image_shape)` ¶

Packed real-spectrum shape for 2D image shape under last-axis rFFT.

`volume_shape_to_half_volume_shape(volume_shape)` ¶

Packed real-spectrum shape for 3D volume shape under last-axis rFFT.

`full_image_to_half_image(image, image_shape)` ¶

Map centered full 2D spectrum to packed Hermitian representation.

Accepts either trailing grid shape (..., *image_shape) or flattened trailing axis (..., prod(image_shape)). Returns in the same style.

`half_image_to_full_image(half_image, image_shape)` ¶

Map packed Hermitian 2D spectrum to centered full complex spectrum.

Uses index-based Hermitian conjugation: non-redundant columns are placed at their packed positions, redundant columns are filled as conj(half[(-i0)%H, |kx|]).

IMPORTANT: Do NOT change this to an FFT-based round-trip (e.g. irfft2 → fft2). This must stay consistent with half_volume_to_full_volume's approach so that VJPs match the CUDA kernel's Hermitian fold scatter.

Accepts either trailing packed grid shape (..., *image_shape_to_half_image_shape(image_shape)) or flattened trailing axis (..., prod(half_shape)). Returns in the same style.

`full_volume_to_half_volume(volume, volume_shape)` ¶

Map centered full 3D spectrum to packed Hermitian representation.

Accepts either trailing grid shape (..., *volume_shape) or flattened trailing axis (..., prod(volume_shape)). Returns in the same style.

`half_volume_to_full_volume(half_volume, volume_shape)` ¶

Map packed Hermitian 3D spectrum to centered full complex spectrum.

Uses index-based Hermitian conjugation: non-redundant columns are placed at their packed positions, redundant columns (negative last-axis frequencies) are filled as conj(half[(-i0)%N0, (-i1)%N1, |kz|]).

This approach has a simple VJP that matches per-voxel Hermitian folding in the CUDA backproject kernel (unlike the FFT round-trip which distributes gradients through the FFT chain differently).

IMPORTANT: Do NOT change this to an FFT-based round-trip (e.g. irfft3 → fft3). The VJP of an FFT round-trip distributes gradients differently than per-voxel Hermitian folding, breaking the CUDA half-volume backproject kernel's correctness. The CUDA kernel's Hermitian fold scatter is specifically designed to be the correct adjoint of THIS index-based expand.

Accepts either trailing packed grid shape (..., *volume_shape_to_half_volume_shape(volume_shape)) or flattened trailing axis (..., prod(half_shape)). Returns in the same style.

`get_dft2_real(img, norm=DEFAULT_FFT_NORM)` ¶

Centered 2D FFT for real-valued inputs using Hermitian packing.

Output shape is (..., H, W//2 + 1) with only the non-redundant last-axis frequencies stored.

`get_idft2_real(img, image_shape=None, norm=DEFAULT_FFT_NORM)` ¶

Inverse of get_dft2_real returning a real-valued image.

`get_dft3_real(img, norm=DEFAULT_FFT_NORM, axes=(-3, -2, -1))` ¶

Centered 3D FFT for real-valued inputs using Hermitian packing.

The packed axis is the last transform axis in axes.

`get_idft3_real(img, volume_shape=None, norm=DEFAULT_FFT_NORM, axes=(-3, -2, -1))` ¶

Inverse of get_dft3_real returning a real-valued volume.

cubic_interpolation¶

Cubic spline interpolation on 3-D grids.

`recovar.core.cubic_interpolation` ¶

JAX-based cubic spline interpolation for multi-dimensional arrays.

Provides cubic spline interpolation functionality with precomputed coefficients.

Uses periodic boundary conditions (circulant system solved via FFT), which preserve Hermitian symmetry of the input: if V[k] = conj(V[-k]) then C[k] = conj(C[-k]). Output shape equals input shape (no boundary padding).

The periodic circulant system [1, 4, 1] is solved as: C = ifft(fft(V) / eigenvalues) where eigenvalues_k = 4 + 2cos(2pi*k/N).

`calculate_spline_coefficients(data)` ¶

Calculate periodic cubic spline coefficients for an array.

Solves the periodic circulant system [1, 4, 1] along each dimension using FFT. Output has the same shape as input (no boundary padding).

Preserves Hermitian symmetry: if data[k] = conj(data[-k]) then the coefficients satisfy the same property.

Parameters:

Name	Type	Description	Default
`data`	`Array`	Input array (any number of dimensions).	required

Returns:

Type	Description
`Array`	Spline coefficients with same shape as input.

`interpolate_with_spline(spline_coeffs, coords, boundary_mode='wrap', fill_value=0.0)` ¶

Interpolate using cubic spline coefficients.

Parameters:

Name	Type	Description	Default
`spline_coeffs`	`ArrayLike`	Precomputed spline coefficients from calculate_spline_coefficients	required
`coords`	`Sequence[ArrayLike]`	Sequence of coordinate arrays, one per dimension	required
`boundary_mode`	`str`	Boundary handling mode ('wrap' for periodic)	`'wrap'`
`fill_value`	`ArrayLike`	Fill value for out-of-bounds when boundary_mode='fill'	`0.0`

Returns:

Type	Description
`Array`	Interpolated values at the specified coordinates

`map_coordinates(input, coordinates, order, mode='wrap', cval=0.0)` ¶

Cubic spline interpolation compatible with scipy/cryojax API.

Only supports order=3 (cubic spline interpolation). Coordinates are shifted by -1 for periodic convention.

Parameters:

Name	Description	Default
`input`	Input array to interpolate	required
`coordinates`	Sequence of coordinate arrays	required
`order`	Interpolation order (must be 3)	required
`mode`	Boundary mode (default 'wrap' for periodic)	`'wrap'`
`cval`	Fill value for out-of-bounds	`0.0`

Returns:

Type	Description
	Interpolated values

`map_coordinates_with_cubic_spline(coefficients, coordinates, mode='wrap', cval=0.0)` ¶

Interpolate using precomputed cubic spline coefficients.

Coordinates are shifted by -1 for periodic convention.

Parameters:

Name	Description	Default
`coefficients`	Precomputed spline coefficients (ndim-dimensional).	required
`coordinates`	Either a sequence of ndim coordinate arrays, or a single array of shape `(ndim, ...)` (scipy convention).	required
`mode`	Boundary mode (default 'wrap' for periodic).	`'wrap'`
`cval`	Fill value for out-of-bounds.	`0.0`

padding¶

Real-space zero-padding and unpadding for Fourier oversampling.

`recovar.core.padding` ¶

Real-space zero-padding and unpadding for Fourier oversampling.

`padded_rfft(images, image_size, padding)` ¶

Pad real-space images and compute real FFT → half-spectrum output.

Like :func:padded_dft but uses rfft2, producing flattened half-spectrum images of shape (n_images, H * (W // 2 + 1)).

recovar.core¶

configs¶

recovar.core.configs ¶

ForwardModelConfig ¶

base_volume_shape property ¶

compute_ctf(ctf_params, *, half_image=False) ¶

Parameters¶

compute_ctf_half(ctf_params) ¶

compute_ctf_at_shape(ctf_params, image_shape) ¶

replace(**kwargs) ¶

from_dataset(cryo, disc_type='linear_interp', process_fn=None, upsampling_factor=None) classmethod ¶

Parameters¶

ModelState ¶

CovColumnOpts ¶

CovarianceOpts ¶

EmbeddingOpts ¶

ctf¶

recovar.core.ctf ¶

CTFParamIndex ¶

CTFMode ¶

CTFEvaluator ¶

Parameters¶

__call__(ctf_params, image_shape, voxel_size, *, half_image=False) ¶

Parameters¶

as_ctf_evaluator(fn_or_evaluator) ¶

evaluate_ctf(freqs, ctf_params) ¶

Parameters¶

Returns¶

compute_antialiased_ctf_squared(ctf_fn, ctf_params, image_shape, voxel_size, *, half_image=False) ¶

geometry¶

recovar.core.geometry ¶

get_gridpoint_coords(rotation_matrix, image_shape, volume_shape) ¶

get_unrotated_half_plane_grid_points(image_shape, three_d_upsampling_factor=1) ¶

translate_images(image, translation, image_shape, *, half_image=False) ¶

forward¶

recovar.core.forward ¶

forward_model(config, volume, ctf_params, rotation_matrices, skip_ctf=False, half_image=False, half_volume=False) ¶

Parameters¶

forward_model_and_adjoint(config, volume, ctf_params, rotation_matrices, skip_ctf=False) ¶

adjoint_forward_model(config, slices, ctf_params, rotation_matrices, skip_ctf=False, volume=None, half_image=False, half_volume=False) ¶

Parameters¶

compute_AtAv(config, volume, ctf_params, rotation_matrices, noise_variance, skip_ctf=False) ¶

slicing¶

recovar.core.slicing ¶

slice_volume(volume, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO) ¶

Parameters¶

batch_slice_volume(volumes, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO) ¶

Parameters¶

adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volume=None, half_image=False, half_volume=False, max_r=_AUTO) ¶

Parameters¶

batch_adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volumes=None, half_image=False, half_volume=False, max_r=_AUTO) ¶

Parameters¶

precompute_cubic_coefficients(volume, volume_shape) ¶

precompute_cubic_coefficients_half(volume, volume_shape) ¶

slice_from_cubic_coefficients(coeffs, rotation_matrices, image_shape, volume_shape, half_image=False) ¶

Parameters¶

indexing¶

recovar.core.indexing ¶

mask¶

recovar.core.mask ¶

masking_options(volume_mask_option, means, volume_shape, dtype_real=np.float32, mask_dilation_iter=0, keep_input_mask=False, dilated_mask_dilations_iter=None) ¶

make_mask(volume, *, threshold='auto', lowpass_sigma=None, extend=None, soft_edge=3, cleanup=True) ¶

make_mask_from_half_maps(halfmap1, halfmap2, smax=3, method='auto', **kwargs) ¶

make_mask_from_gt(gt_map, smax=3, iter=10, from_ft=True) ¶

make_union_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3) ¶

soften_volume_mask(binary_volume_mask, kern_rad=3) ¶

get_radial_mask(shape, radius=None) ¶

raised_cosine_mask(volume_shape, radius, radius_p, offset) ¶

soft_mask_outside_map(vol, radius=-1, cosine_width=3, Mnoise=None) ¶

window_mask(D, in_rad, out_rad) ¶

smooth_circular_mask(image_size, radius, thickness) ¶

make_soft_edged_kernel(r1, shape) ¶

threshold_map(arr, prob=0.99, dthresh=None) ¶

make_moving_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3) ¶

linalg¶

recovar.core.linalg ¶

thin_svd_in_blocks(X, np=np, memory_to_use=5, epsilon=1e-08, n_components=-1) ¶

thin_svd(X, np=np, epsilon=1e-08) ¶

randomized_svd(A, n_pcs=200) ¶

broadcast_dot(x, y) ¶

`recovar.core.configs` ¶

`ForwardModelConfig` ¶

`base_volume_shape` `property` ¶

`compute_ctf(ctf_params, *, half_image=False)` ¶

`compute_ctf_half(ctf_params)` ¶

`compute_ctf_at_shape(ctf_params, image_shape)` ¶

`replace(**kwargs)` ¶

`from_dataset(cryo, disc_type='linear_interp', process_fn=None, upsampling_factor=None)` `classmethod` ¶

`ModelState` ¶

`CovColumnOpts` ¶

`CovarianceOpts` ¶

`EmbeddingOpts` ¶

`recovar.core.ctf` ¶

`CTFParamIndex` ¶

`CTFMode` ¶

`CTFEvaluator` ¶

`call(ctf_params, image_shape, voxel_size, *, half_image=False)` ¶

`as_ctf_evaluator(fn_or_evaluator)` ¶

`evaluate_ctf(freqs, ctf_params)` ¶

`compute_antialiased_ctf_squared(ctf_fn, ctf_params, image_shape, voxel_size, *, half_image=False)` ¶

`recovar.core.geometry` ¶

`get_gridpoint_coords(rotation_matrix, image_shape, volume_shape)` ¶

`get_unrotated_half_plane_grid_points(image_shape, three_d_upsampling_factor=1)` ¶

`translate_images(image, translation, image_shape, *, half_image=False)` ¶

`recovar.core.forward` ¶

`forward_model(config, volume, ctf_params, rotation_matrices, skip_ctf=False, half_image=False, half_volume=False)` ¶

`forward_model_and_adjoint(config, volume, ctf_params, rotation_matrices, skip_ctf=False)` ¶

`adjoint_forward_model(config, slices, ctf_params, rotation_matrices, skip_ctf=False, volume=None, half_image=False, half_volume=False)` ¶

`compute_AtAv(config, volume, ctf_params, rotation_matrices, noise_variance, skip_ctf=False)` ¶

`recovar.core.slicing` ¶

`slice_volume(volume, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO)` ¶

`batch_slice_volume(volumes, rotation_matrices, image_shape, volume_shape, disc_type, half_volume=False, half_image=False, max_r=_AUTO)` ¶

`adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volume=None, half_image=False, half_volume=False, max_r=_AUTO)` ¶

`batch_adjoint_slice_volume(slices, rotation_matrices, image_shape, volume_shape, disc_type, volumes=None, half_image=False, half_volume=False, max_r=_AUTO)` ¶

`precompute_cubic_coefficients(volume, volume_shape)` ¶

`precompute_cubic_coefficients_half(volume, volume_shape)` ¶

`slice_from_cubic_coefficients(coeffs, rotation_matrices, image_shape, volume_shape, half_image=False)` ¶

`recovar.core.indexing` ¶

`recovar.core.mask` ¶

`masking_options(volume_mask_option, means, volume_shape, dtype_real=np.float32, mask_dilation_iter=0, keep_input_mask=False, dilated_mask_dilations_iter=None)` ¶

`make_mask(volume, *, threshold='auto', lowpass_sigma=None, extend=None, soft_edge=3, cleanup=True)` ¶

`make_mask_from_half_maps(halfmap1, halfmap2, smax=3, method='auto', **kwargs)` ¶

`make_mask_from_gt(gt_map, smax=3, iter=10, from_ft=True)` ¶

`make_union_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3)` ¶

`soften_volume_mask(binary_volume_mask, kern_rad=3)` ¶

`get_radial_mask(shape, radius=None)` ¶

`raised_cosine_mask(volume_shape, radius, radius_p, offset)` ¶

`soft_mask_outside_map(vol, radius=-1, cosine_width=3, Mnoise=None)` ¶

`window_mask(D, in_rad, out_rad)` ¶

`smooth_circular_mask(image_size, radius, thickness)` ¶

`make_soft_edged_kernel(r1, shape)` ¶

`threshold_map(arr, prob=0.99, dthresh=None)` ¶

`make_moving_gt_mask(gt_volumes_real, volume_shape, smax=3, iter=1, dilation_iters=None, kern_rad=3)` ¶

`recovar.core.linalg` ¶

`thin_svd_in_blocks(X, np=np, memory_to_use=5, epsilon=1e-08, n_components=-1)` ¶

`thin_svd(X, np=np, epsilon=1e-08)` ¶

`randomized_svd(A, n_pcs=200)` ¶

`broadcast_dot(x, y)` ¶

`broadcast_outer(x, y)` ¶

`inner_product(x, y)` ¶

`batch_inner_product(x, y)` ¶

`half_spectrum_last_axis_weights(last_axis_size, dtype=jnp.float32)` ¶

`rfft2_hermitian_weights(image_shape, dtype=jnp.float32)` ¶

`half_spectrum_inner_product(x_half, y_half, full_shape)` ¶

`batch_half_spectrum_inner_product(x_half, y_half, full_shape)` ¶

`recovar.core.fourier_transform_utils` ¶

`get_1d_frequency_grid_rfft(n, voxel_size=1, scaled=False)` ¶

`get_k_coordinate_of_each_pixel_real(image_shape, voxel_size, scaled=True)` ¶

`get_k_coordinate_of_each_pixel_3d_real(image_shape, voxel_size, scaled=True)` ¶

`get_grid_of_radial_distances_real(image_shape, voxel_size=1, scaled=False, frequency_shift=0, rounded=True)` ¶

`get_real_fft_packed_shape(shape)` ¶

`get_real_fft_packed_last_axis_indices(n)` ¶

`get_k_coordinate_of_each_pixel_half(image_shape, voxel_size, scaled=True)` ¶

`get_shifted_conjugate_partner_indices(n)` ¶

`get_real_fft_memory_saving_ratio(shape)` ¶

`image_shape_to_half_image_shape(image_shape)` ¶

`volume_shape_to_half_volume_shape(volume_shape)` ¶

`full_image_to_half_image(image, image_shape)` ¶

`half_image_to_full_image(half_image, image_shape)` ¶

`full_volume_to_half_volume(volume, volume_shape)` ¶