utils

Utility functions shared between pydantic models.

calculate_ctf_filter_stack(template_shape, optics_group, defocus_offsets, pixel_size_offsets)

Calculate searched CTF filter values for a given shape and optics group.

Parameters:

Name	Type	Description	Default
template_shape	tuple[int, int]	Desired output shape for the filter, in real space.	required
optics_group	OpticsGroup	OpticsGroup object containing the optics defining the CTF parameters.	required
defocus_offsets	Tensor	Tensor of defocus offsets to search over, in Angstroms.	required
pixel_size_offsets	Tensor	Tensor of pixel size offsets to search over, in Angstroms.	required

Returns:

Type	Description
Tensor	Tensor of CTF filter values for the specified shape and optics group. Will have shape (num_pixel_sizes, num_defocus_offsets, h, w // 2 + 1)

Source code in src/leopard_em/pydantic_models/utils.py

def calculate_ctf_filter_stack(
    template_shape: tuple[int, int],
    optics_group: "OpticsGroup",
    defocus_offsets: torch.Tensor,  # in Angstrom, relative
    pixel_size_offsets: torch.Tensor,  # in Angstrom, relative
) -> torch.Tensor:
    """Calculate searched CTF filter values for a given shape and optics group.

    Parameters
    ----------
    template_shape : tuple[int, int]
        Desired output shape for the filter, in real space.
    optics_group : OpticsGroup
        OpticsGroup object containing the optics defining the CTF parameters.
    defocus_offsets : torch.Tensor
        Tensor of defocus offsets to search over, in Angstroms.
    pixel_size_offsets : torch.Tensor
        Tensor of pixel size offsets to search over, in Angstroms.

    Returns
    -------
    torch.Tensor
        Tensor of CTF filter values for the specified shape and optics group. Will have
        shape (num_pixel_sizes, num_defocus_offsets, h, w // 2 + 1)
    """
    return calculate_ctf_filter_stack_full_args(
        template_shape,
        optics_group.defocus_u,
        optics_group.defocus_v,
        defocus_offsets,
        pixel_size_offsets,
        astigmatism_angle=optics_group.astigmatism_angle,
        voltage=optics_group.voltage,
        spherical_aberration=optics_group.spherical_aberration,
        amplitude_contrast_ratio=optics_group.amplitude_contrast_ratio,
        ctf_B_factor=optics_group.ctf_B_factor,
        phase_shift=optics_group.phase_shift,
        pixel_size=optics_group.pixel_size,
    )

calculate_ctf_filter_stack_full_args(template_shape, defocus_u, defocus_v, defocus_offsets, pixel_size_offsets, **kwargs)

Calculate a CTF filter stack for a given set of parameters and search offsets.

Parameters:

Name	Type	Description	Default
template_shape	tuple[int, int]	Desired output shape for the filter, in real space.	required
defocus_u	float	Defocus along the major axis, in Angstroms.	required
defocus_v	float	Defocus along the minor axis, in Angstroms.	required
defocus_offsets	Tensor	Tensor of defocus offsets to search over, in Angstroms.	required
pixel_size_offsets	Tensor	Tensor of pixel size offsets to search over, in Angstroms.	required
**kwargs	Any	Additional keyword to pass to the calculate_ctf_2d function.	{}

Returns:

Type	Description
Tensor	Tensor of CTF filter values for the specified shape and parameters. Will have shape (num_pixel_sizes, num_defocus_offsets, h, w // 2 + 1)
# Raises
# ------
# ValueError
# If not all the required parameters are passed as additional keyword arguments.

Source code in src/leopard_em/pydantic_models/utils.py

def calculate_ctf_filter_stack_full_args(
    template_shape: tuple[int, int],
    defocus_u: float,  # in Angstrom
    defocus_v: float,  # in Angstrom
    defocus_offsets: torch.Tensor,  # in Angstrom, relative
    pixel_size_offsets: torch.Tensor,  # in Angstrom, relative
    **kwargs: Any,
) -> torch.Tensor:
    """Calculate a CTF filter stack for a given set of parameters and search offsets.

    Parameters
    ----------
    template_shape : tuple[int, int]
        Desired output shape for the filter, in real space.
    defocus_u : float
        Defocus along the major axis, in Angstroms.
    defocus_v : float
        Defocus along the minor axis, in Angstroms.
    defocus_offsets : torch.Tensor
        Tensor of defocus offsets to search over, in Angstroms.
    pixel_size_offsets : torch.Tensor
        Tensor of pixel size offsets to search over, in Angstroms.
    **kwargs
        Additional keyword to pass to the calculate_ctf_2d function.

    Returns
    -------
    torch.Tensor
        Tensor of CTF filter values for the specified shape and parameters. Will have
        shape (num_pixel_sizes, num_defocus_offsets, h, w // 2 + 1)

    # Raises
    # ------
    # ValueError
    #     If not all the required parameters are passed as additional keyword arguments.
    """
    # Calculate the defocus values + offsets in terms of Angstrom
    defocus = defocus_offsets + ((defocus_u + defocus_v) / 2)
    astigmatism = abs(defocus_u - defocus_v) / 2

    # The different Cs values to search over as another dimension
    cs_values = get_cs_range(
        pixel_size=kwargs["pixel_size"],
        pixel_size_offsets=pixel_size_offsets,
        cs=kwargs["spherical_aberration"],
    )

    # Ensure defocus and astigmatism have a batch dimension so Cs and defocus can be
    # interleaved correctly
    if defocus.dim() == 0:
        defocus = defocus.unsqueeze(0)

    # Loop over spherical aberrations one at a time and collect results
    ctf_list = []
    for cs_val in cs_values:
        tmp = calculate_ctf_2d(
            defocus=defocus * 1e-4,  # Convert to um from Angstrom
            astigmatism=astigmatism * 1e-4,  # Convert to um from Angstrom
            astigmatism_angle=kwargs["astigmatism_angle"],
            voltage=kwargs["voltage"],
            spherical_aberration=cs_val,
            amplitude_contrast=kwargs["amplitude_contrast_ratio"],
            b_factor=kwargs["ctf_B_factor"],
            phase_shift=kwargs["phase_shift"],
            pixel_size=kwargs["pixel_size"],
            image_shape=template_shape,
            rfft=True,
            fftshift=False,
        )
        ctf_list.append(tmp)

    ctf = torch.stack(ctf_list, dim=0)

    return ctf

cs_to_pixel_size(cs_vals, nominal_pixel_size, nominal_cs=2.7)

Convert Cs values to pixel sizes.

Parameters:

Name	Type	Description	Default
cs_vals	Tensor	The Cs (spherical aberration) values.	required
nominal_pixel_size	float	The nominal pixel size.	required
nominal_cs	float	The nominal Cs value, by default 2.7.	2.7

Returns:

Type	Description
Tensor	The pixel sizes.

Source code in src/leopard_em/pydantic_models/utils.py

def cs_to_pixel_size(
    cs_vals: torch.Tensor,
    nominal_pixel_size: float,
    nominal_cs: float = 2.7,
) -> torch.Tensor:
    """Convert Cs values to pixel sizes.

    Parameters
    ----------
    cs_vals : torch.Tensor
        The Cs (spherical aberration) values.
    nominal_pixel_size : float
        The nominal pixel size.
    nominal_cs : float, optional
        The nominal Cs value, by default 2.7.

    Returns
    -------
    torch.Tensor
        The pixel sizes.
    """
    pixel_size = torch.pow(nominal_cs / cs_vals, 0.25) * nominal_pixel_size
    return pixel_size

get_cs_range(pixel_size, pixel_size_offsets, cs=2.7)

Get the Cs values for a range of pixel sizes.

Parameters:

Name	Type	Description	Default
pixel_size	float	The nominal pixel size.	required
pixel_size_offsets	Tensor	The pixel size offsets.	required
cs	float	The Cs (spherical aberration) value, by default 2.7.	2.7

Returns:

Type	Description
Tensor	The Cs values for the range of pixel sizes.

Source code in src/leopard_em/pydantic_models/utils.py

def get_cs_range(
    pixel_size: float,
    pixel_size_offsets: torch.Tensor,
    cs: float = 2.7,
) -> torch.Tensor:
    """Get the Cs values for a  range of pixel sizes.

    Parameters
    ----------
    pixel_size : float
        The nominal pixel size.
    pixel_size_offsets : torch.Tensor
        The pixel size offsets.
    cs : float, optional
        The Cs (spherical aberration) value, by default 2.7.

    Returns
    -------
    torch.Tensor
        The Cs values for the range of pixel sizes.
    """
    pixel_sizes = pixel_size + pixel_size_offsets
    cs_values = cs / torch.pow(pixel_sizes / pixel_size, 4)
    return cs_values

get_search_tensors(min_val, max_val, step_size, skip_enforce_zero=False)

Get the search tensors (pixel or defocus) for a given range and step size.

Parameters:

Name	Type	Description	Default
min_val	float	The minimum value.	required
max_val	float	The maximum value.	required
step_size	float	The step size.	required
skip_enforce_zero	bool	Whether to skip enforcing a zero value, by default False.	False

Returns:

Type	Description
tensor	The search tensors.

Source code in src/leopard_em/pydantic_models/utils.py

def get_search_tensors(
    min_val: float,
    max_val: float,
    step_size: float,
    skip_enforce_zero: bool = False,
) -> torch.tensor:
    """Get the search tensors (pixel or defocus) for a given range and step size.

    Parameters
    ----------
    min_val : float
        The minimum value.
    max_val : float
        The maximum value.
    step_size : float
        The step size.
    skip_enforce_zero : bool, optional
        Whether to skip enforcing a zero value, by default False.

    Returns
    -------
    torch.tensor
        The search tensors.
    """
    vals = torch.arange(
        min_val,
        max_val + step_size,
        step_size,
        dtype=torch.float32,
    )

    if abs(torch.min(torch.abs(vals))) > 1e-6 and not skip_enforce_zero:
        vals = torch.cat([vals, torch.tensor([0.0])])
        # Re-sort pixel sizes
        vals = torch.sort(vals)[0]

    return vals

preprocess_image(image_rfft, cumulative_fourier_filters, bandpass_filter)

Preprocesses and normalizes the image based on the given filters.

Parameters:

Name	Type	Description	Default
image_rfft	Tensor	The real Fourier-transformed image (unshifted).	required
cumulative_fourier_filters	Tensor	The cumulative Fourier filters. Multiplication of the whitening filter, phase randomization filter, bandpass filter, and arbitrary curve filter.	required
bandpass_filter	Tensor	The bandpass filter used for the image. Used for dimensionality normalization.	required

Returns:

Type	Description
Tensor	Preprocessed and normalized image in Fourier space

Source code in src/leopard_em/pydantic_models/utils.py

def preprocess_image(
    image_rfft: torch.Tensor,
    cumulative_fourier_filters: torch.Tensor,
    bandpass_filter: torch.Tensor,
) -> torch.Tensor:
    """Preprocesses and normalizes the image based on the given filters.

    Parameters
    ----------
    image_rfft : torch.Tensor
        The real Fourier-transformed image (unshifted).
    cumulative_fourier_filters : torch.Tensor
        The cumulative Fourier filters. Multiplication of the whitening filter, phase
        randomization filter, bandpass filter, and arbitrary curve filter.
    bandpass_filter : torch.Tensor
        The bandpass filter used for the image. Used for dimensionality normalization.

    Returns
    -------
    torch.Tensor
        Preprocessed and normalized image in Fourier space
    """
    image_rfft = image_rfft * cumulative_fourier_filters

    # Normalize the image after filtering
    squared_image_rfft = torch.abs(image_rfft) ** 2
    squared_sum = torch.sum(squared_image_rfft, dim=(-2, -1), keepdim=True)
    squared_sum += torch.sum(
        squared_image_rfft[..., :, 1:-1], dim=(-2, -1), keepdim=True
    )
    image_rfft /= torch.sqrt(squared_sum)

    # NOTE: For two Gaussian random variables in d-dimensional space --  A and B --
    # each with mean 0 and variance 1 their correlation will have on average a
    # variance of d.
    # NOTE: Since we have the variance of the image and template projections each at
    # 1, we need to multiply the image by the square root of the number of pixels
    # so the cross-correlograms have a variance of 1 and not d.
    # NOTE: When applying the Fourier filters to the image and template, any
    # elements that get set to zero effectively reduce the dimensionality of our
    # cross-correlation. Therefore, instead of multiplying by the number of pixels,
    # we need to multiply tby the effective number of pixels that are non-zero.
    # Below, we calculate the dimensionality of our cross-correlation and divide
    # by the square root of that number to normalize the image.
    dimensionality = bandpass_filter.sum() + bandpass_filter[:, 1:-1].sum()
    image_rfft *= dimensionality**0.5

    return image_rfft

setup_images_filters_particle_stack(particle_stack, preprocessing_filters, template)

Extract and preprocess particle images and calculate filters.

This function extracts particle images from a particle stack, performs FFT, applies filters, and prepares the template for further processing.

Parameters:

Name	Type	Description	Default
particle_stack	ParticleStack	The particle stack containing images to process.	required
preprocessing_filters	PreprocessingFilters	Filters to apply to the particle images.	required
template	Tensor	The 3D template volume.	required

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	A tuple containing: - particle_images_dft: The particle images in Fourier space - template_dft: The Fourier transformed template - projective_filters: Filters applied to the template

Source code in src/leopard_em/pydantic_models/utils.py

def setup_images_filters_particle_stack(
    particle_stack: "ParticleStack",
    preprocessing_filters: "PreprocessingFilters",
    template: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Extract and preprocess particle images and calculate filters.

    This function extracts particle images from a particle stack, performs FFT,
    applies filters, and prepares the template for further processing.

    Parameters
    ----------
    particle_stack : ParticleStack
        The particle stack containing images to process.
    preprocessing_filters : PreprocessingFilters
        Filters to apply to the particle images.
    template : torch.Tensor
        The 3D template volume.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        A tuple containing:
        - particle_images_dft: The particle images in Fourier space
        - template_dft: The Fourier transformed template
        - projective_filters: Filters applied to the template
    """
    # Extract out the regions of interest (particles) based on the particle stack
    particle_images = particle_stack.construct_image_stack(
        pos_reference="center",
        padding_value=0.0,
        handle_bounds="pad",
        padding_mode="constant",
    )

    # FFT the particle images
    # pylint: disable=E1102
    particle_images_dft = torch.fft.rfftn(particle_images, dim=(-2, -1))
    particle_images_dft[..., 0, 0] = 0.0 + 0.0j  # Zero out DC component

    bandpass_filter = preprocessing_filters.bandpass_filter.calculate_bandpass_filter(
        particle_images_dft.shape[-2:]
    )

    # Calculate and apply the filters for the particle image stack
    filter_stack = particle_stack.construct_filter_stack(
        preprocessing_filters, output_shape=particle_images_dft.shape[-2:]
    )

    particle_images_dft = preprocess_image(
        image_rfft=particle_images_dft,
        cumulative_fourier_filters=filter_stack,
        bandpass_filter=bandpass_filter,
    )

    # Calculate the filters applied to each template (besides CTF)
    projective_filters = particle_stack.construct_filter_stack(
        preprocessing_filters,
        output_shape=(template.shape[-2], template.shape[-1] // 2 + 1),
    )

    template_dft = volume_to_rfft_fourier_slice(template)

    return (
        particle_images_dft,
        template_dft,
        projective_filters,
    )

setup_particle_backend_kwargs(particle_stack, template, preprocessing_filters, euler_angles, euler_angle_offsets, defocus_offsets, pixel_size_offsets, device_list)

Create common kwargs dictionary for template backend functions.

This function extracts the common code between RefineTemplateManager and OptimizeTemplateManager's make_backend_core_function_kwargs methods.

Parameters:

Name	Type	Description	Default
particle_stack	ParticleStack	The particle stack containing images to process.	required
template	Tensor	The 3D template volume.	required
preprocessing_filters	PreprocessingFilters	Filters to apply to the particle images.	required
euler_angles	Tensor	The set of Euler angles to use.	required
euler_angle_offsets	Tensor	The relative Euler angle offsets to search over.	required
defocus_offsets	Tensor	The relative defocus values to search over.	required
pixel_size_offsets	Tensor	The relative pixel size values to search over.	required
device_list	list	List of computational devices to use.	required

Returns:

Type	Description
dict[str, Any]	Dictionary of keyword arguments for backend functions.

Source code in src/leopard_em/pydantic_models/utils.py

def setup_particle_backend_kwargs(
    particle_stack: "ParticleStack",
    template: torch.Tensor,
    preprocessing_filters: "PreprocessingFilters",
    euler_angles: torch.Tensor,
    euler_angle_offsets: torch.Tensor,
    defocus_offsets: torch.Tensor,
    pixel_size_offsets: torch.Tensor,
    device_list: list,
) -> dict[str, Any]:
    """Create common kwargs dictionary for template backend functions.

    This function extracts the common code between RefineTemplateManager and
    OptimizeTemplateManager's make_backend_core_function_kwargs methods.

    Parameters
    ----------
    particle_stack : ParticleStack
        The particle stack containing images to process.
    template : torch.Tensor
        The 3D template volume.
    preprocessing_filters : PreprocessingFilters
        Filters to apply to the particle images.
    euler_angles : torch.Tensor
        The set of Euler angles to use.
    euler_angle_offsets : torch.Tensor
        The relative Euler angle offsets to search over.
    defocus_offsets : torch.Tensor
        The relative defocus values to search over.
    pixel_size_offsets : torch.Tensor
        The relative pixel size values to search over.
    device_list : list
        List of computational devices to use.

    Returns
    -------
    dict[str, Any]
        Dictionary of keyword arguments for backend functions.
    """
    # Get correlation statistics
    corr_mean_stack = particle_stack.construct_cropped_statistic_stack(
        "correlation_average",
    )
    corr_std_stack = (
        particle_stack.construct_cropped_statistic_stack(
            stat="correlation_variance",
            pos_reference="center",
            handle_bounds="pad",
            padding_mode="constant",
            padding_value=1e10,  # large to avoid out of bound pixels having inf z-score
        )
        ** 0.5
    )  # var to std

    # Extract and preprocess images and filters
    (
        particle_images_dft,
        template_dft,
        projective_filters,
    ) = setup_images_filters_particle_stack(
        particle_stack, preprocessing_filters, template
    )

    # The best defocus values for each particle (+ astigmatism)
    defocus_u, defocus_v = particle_stack.get_absolute_defocus()
    defocus_angle = torch.tensor(particle_stack["astigmatism_angle"])

    ctf_kwargs = _setup_ctf_kwargs_from_particle_stack(
        particle_stack, (template.shape[-2], template.shape[-1])
    )

    return {
        "particle_stack_dft": particle_images_dft,
        "template_dft": template_dft,
        "euler_angles": euler_angles,
        "euler_angle_offsets": euler_angle_offsets,
        "defocus_u": defocus_u,
        "defocus_v": defocus_v,
        "defocus_angle": defocus_angle,
        "defocus_offsets": defocus_offsets,
        "pixel_size_offsets": pixel_size_offsets,
        "corr_mean": corr_mean_stack,
        "corr_std": corr_std_stack,
        "ctf_kwargs": ctf_kwargs,
        "projective_filters": projective_filters,
        "device": device_list,
    }

volume_to_rfft_fourier_slice(volume)

Prepares a 3D volume for Fourier slice extraction.

Parameters:

Name	Type	Description	Default
volume	Tensor	The input volume.	required

Returns:

Type	Description
Tensor	The prepared volume in Fourier space ready for slice extraction.

Source code in src/leopard_em/pydantic_models/utils.py

def volume_to_rfft_fourier_slice(volume: torch.Tensor) -> torch.Tensor:
    """Prepares a 3D volume for Fourier slice extraction.

    Parameters
    ----------
    volume : torch.Tensor
        The input volume.

    Returns
    -------
    torch.Tensor
        The prepared volume in Fourier space ready for slice extraction.
    """
    assert volume.dim() == 3, "Volume must be 3D"

    # NOTE: There is an extra FFTshift step before the RFFT since, for some reason,
    # omitting this step will cause a 180 degree phase shift on odd (i, j, k)
    # structure factors in the Fourier domain. This just requires an extra
    # IFFTshift after converting a slice back to real-space (handled already).
    volume = torch.fft.fftshift(volume, dim=(0, 1, 2))  # pylint: disable=E1102
    volume_rfft = torch.fft.rfftn(volume, dim=(0, 1, 2))  # pylint: disable=E1102
    volume_rfft = torch.fft.fftshift(volume_rfft, dim=(0, 1))  # pylint: disable=E1102

    return volume_rfft