FSC (needs to be checked)

fsc.py

from pathlib import Path

import einops
import mrcfile
import numpy as np
import torch

import typer

cli = typer.Typer(add_completion=False)


def calculate_fsc(a: torch.Tensor, b: torch.Tensor):
    """Calculate the Fourier shell correlation between two cubic volume."""
    # specify position of the DC component of the rfft
    rfftn_dc_idx = torch.div(torch.Tensor(tuple(a.shape)), 2,
                             rounding_mode='floor')
    rfftn_dc_idx[-1] = 0

    # calculate DFTs of volume a and b
    a, b = torch.fft.rfftn(a), torch.fft.rfftn(b)
    a, b = torch.fft.fftshift(a, dim=(-3, -2)), torch.fft.fftshift(a,
                                                                   dim=(-3, -2))
    n_fsc = a.shape[-1]

    # calculate distance from DC component for each fourier coefficient
    array_indices = torch.Tensor(np.indices(a.shape))  # (c, (z), y, x)
    array_indices = einops.rearrange(array_indices, 'c ... -> ... c')
    distances = torch.linalg.norm(array_indices - rfftn_dc_idx, dim=-1)

    # linearise array and distances then sort on distance to enable splitting
    # into shells in one pass
    a, b = torch.flatten(a), torch.flatten(b)
    distances = torch.flatten(distances)
    distances, distances_sorted_idx = torch.sort(distances, descending=False)

    # find indices for fourier features in each shell
    split_at_idx = torch.searchsorted(distances, torch.arange(n_fsc)[1:])
    shell_vector_idx = torch.tensor_split(distances_sorted_idx, split_at_idx)

    # extract shells as separate arrays
    shells_a = [a[idx] for idx in shell_vector_idx]
    shells_b = [b[idx] for idx in shell_vector_idx]

    # calculate the correlation in each shell
    fsc = [
        torch.dot(ai, torch.conj(bi)) / (
            torch.linalg.norm(ai) * torch.linalg.norm(bi))
        for ai, bi
        in zip(shells_a, shells_b)
    ]

    return torch.real(torch.tensor(fsc))


@cli.command(no_args_is_help=True)
def fsc_between_volume_files(volume_a: Path, volume_b: Path):
    a, b = mrcfile.read(volume_a), mrcfile.read(volume_b)
    a, b = torch.tensor(a), torch.tensor(b)
    fsc = calculate_fsc(a, b)
    for f in fsc:
        typer.echo(f'{f:.3f}')


if __name__ == '__main__':
    cli()