diff --git a/pyproject.toml b/pyproject.toml index 3db2a2d..6947df8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,7 @@ ttfsc = "ttfsc:cli" # https://docs.astral.sh/ruff [tool.ruff] -line-length = 88 +line-length = 125 target-version = "py38" src = ["src"] @@ -112,7 +112,7 @@ skip-magic-trailing-comma = false # default is false # https://mypy.readthedocs.io/en/stable/config_file.html [tool.mypy] files = "src/**/" -strict = true +strict = false disallow_any_generics = false disallow_subclassing_any = false show_error_codes = true diff --git a/src/ttfsc/__init__.py b/src/ttfsc/__init__.py index 8235d89..3ac582e 100644 --- a/src/ttfsc/__init__.py +++ b/src/ttfsc/__init__.py @@ -1,4 +1,4 @@ -"""Use Fourier shell correlation to estimate the resolution of cryo-EM images and volumes""" +"""Use Fourier shell correlation to estimate the resolution of cryo-EM images and volumes.""" from importlib.metadata import PackageNotFoundError, version @@ -9,4 +9,6 @@ __author__ = "Johannes Elferich" __email__ = "jojotux123@hotmail.com" -from .cli import cli \ No newline at end of file +from ._cli import cli + +__all__ = ["cli"] diff --git a/src/ttfsc/cli.py b/src/ttfsc/_cli.py similarity index 63% rename from src/ttfsc/cli.py rename to src/ttfsc/_cli.py index 9c6f2e2..5eee389 100644 --- a/src/ttfsc/cli.py +++ b/src/ttfsc/_cli.py @@ -1,48 +1,54 @@ from pathlib import Path -from typing import Optional, Annotated +from typing import Annotated, Optional import mrcfile -import typer import torch -from torch_fourier_shell_correlation import fsc +import typer from rich import print as rprint +from torch_fourier_shell_correlation import fsc + +cli = typer.Typer(name="ttfsc", no_args_is_help=True, add_completion=False) -cli = typer.Typer(name='ttctf', no_args_is_help=True, add_completion=False) @cli.command(no_args_is_help=True) def ttfsc_cli( map1: Annotated[Path, typer.Argument(show_default=False)], map2: Annotated[Path, typer.Argument(show_default=False)], - pixel_spacing_angstroms: Annotated[Optional[float], typer.Option('--pixel-spacing-angstroms', show_default=False, help="Pixel spacing in Å/px, taken from header if not set")] = None, - plot: Annotated[bool, typer.Option('--plot')] = True, - plot_with_matplotlib: Annotated[bool, typer.Option('--plot-with-matplotlib')] = False, - fsc_threshold: Annotated[float, typer.Option('--fsc-treshold', show_default=False, help="FSC treshold")] = 0.143, - #mask: Annotated[] -): + pixel_spacing_angstroms: Annotated[ + Optional[float], + typer.Option( + "--pixel-spacing-angstroms", show_default=False, help="Pixel spacing in Å/px, taken from header if not set" + ), + ] = None, + plot: Annotated[bool, typer.Option("--plot")] = True, + plot_with_matplotlib: Annotated[bool, typer.Option("--plot-with-matplotlib")] = False, + fsc_threshold: Annotated[float, typer.Option("--fsc-threshold", show_default=False, help="FSC threshold")] = 0.143, + # mask: Annotated[] +) -> None: with mrcfile.open(map1) as f: map1_tensor = torch.tensor(f.data) if pixel_spacing_angstroms is None: pixel_spacing_angstroms = f.voxel_size.x with mrcfile.open(map2) as f: map2_tensor = torch.tensor(f.data) - - + frequency_pixels = torch.fft.rfftfreq(map1_tensor.shape[0]) resolution_angstroms = (1 / frequency_pixels) * pixel_spacing_angstroms fsc_values = fsc(map1_tensor, map2_tensor) - estimated_resolution_frequency_pixel = float(frequency_pixels[(fsc_values < fsc_threshold).nonzero()[0]-1]) - estimated_resolution_angstrom = float(resolution_angstroms[(fsc_values < fsc_threshold).nonzero()[0]-1]) + estimated_resolution_frequency_pixel = float(frequency_pixels[(fsc_values < fsc_threshold).nonzero()[0] - 1]) + estimated_resolution_angstrom = float(resolution_angstroms[(fsc_values < fsc_threshold).nonzero()[0] - 1]) rprint(f"Estimated resolution using {fsc_threshold} criterion: {estimated_resolution_angstrom:.2f} Å") if plot: - from .plotting import plot_matplotlib, plot_plottile + from ._plotting import plot_matplotlib, plot_plottile + if plot_with_matplotlib: plot_matplotlib( fsc_values=fsc_values, resolution_angstroms=resolution_angstroms, estimated_resolution_angstrom=estimated_resolution_angstrom, - fsc_threshold=fsc_threshold + fsc_threshold=fsc_threshold, ) else: plot_plottile( @@ -50,8 +56,5 @@ def ttfsc_cli( frequency_pixels=frequency_pixels, pixel_spacing_angstroms=pixel_spacing_angstroms, estimated_resolution_frequency_pixel=estimated_resolution_frequency_pixel, - fsc_threshold=fsc_threshold + fsc_threshold=fsc_threshold, ) - - - \ No newline at end of file diff --git a/src/ttfsc/_plotting.py b/src/ttfsc/_plotting.py new file mode 100644 index 0000000..a03ba87 --- /dev/null +++ b/src/ttfsc/_plotting.py @@ -0,0 +1,53 @@ +import torch + + +def plot_matplotlib( + fsc_values: torch.Tensor, resolution_angstroms: torch.Tensor, estimated_resolution_angstrom: float, fsc_threshold: float +) -> None: + from matplotlib import pyplot as plt + + plt.hlines(0, resolution_angstroms[1], resolution_angstroms[-2], "black") + plt.plot(resolution_angstroms, fsc_values, label="FSC (unmasked)") + plt.xlabel("Resolution (Å)") + plt.ylabel("Correlation") + plt.xscale("log") + plt.xlim(resolution_angstroms[1], resolution_angstroms[-2]) + plt.ylim(-0.05, 1.05) + plt.hlines(fsc_threshold, resolution_angstroms[1], estimated_resolution_angstrom, "red", "--") + plt.vlines(estimated_resolution_angstrom, -0.05, fsc_threshold, "red", "--") + plt.legend() + plt.tight_layout() + plt.show() + + +def plot_plottile( + fsc_values: torch.Tensor, + frequency_pixels: torch.Tensor, + pixel_spacing_angstroms: float, + estimated_resolution_frequency_pixel: float, + fsc_threshold: float, +) -> None: + import plotille + + fig = plotille.Figure() + fig.width = 60 + fig.height = 20 + fig.set_x_limits(float(frequency_pixels[1]), float(frequency_pixels[-1])) + fig.set_y_limits(0, 1) + + def resolution_callback(x: float, _: float) -> str: + return f"{(1 / x) * pixel_spacing_angstroms:.2f}" + + fig.x_ticks_fkt = resolution_callback + fig.plot(frequency_pixels[1:].numpy(), fsc_values[1:].numpy(), lc="blue", label="FSC") + + fig.plot( + [float(frequency_pixels[1].numpy()), estimated_resolution_frequency_pixel], + [fsc_threshold, fsc_threshold], + lc="red", + label=" ", + ) + fig.plot( + [estimated_resolution_frequency_pixel, estimated_resolution_frequency_pixel], [0, fsc_threshold], lc="red", label=" " + ) + print(fig.show(legend=True)) diff --git a/src/ttfsc/plotting.py b/src/ttfsc/plotting.py deleted file mode 100644 index 86eb1b7..0000000 --- a/src/ttfsc/plotting.py +++ /dev/null @@ -1,44 +0,0 @@ -import torch - -def plot_matplotlib( - fsc_values: torch.Tensor, - resolution_angstroms: torch.Tensor, - estimated_resolution_angstrom: float, - fsc_threshold: float - ): - from matplotlib import pyplot as plt - plt.hlines(0,resolution_angstroms[1], resolution_angstroms[-2],'black') - plt.plot(resolution_angstroms, fsc_values,label="FSC (unmasked)") - plt.xlabel('Resolution (Å)') - plt.ylabel('Correlation') - plt.xscale('log') - plt.xlim(resolution_angstroms[1], resolution_angstroms[-2]) - plt.ylim(-0.05,1.05) - plt.hlines(fsc_threshold,resolution_angstroms[1],estimated_resolution_angstrom,'red','--') - plt.vlines(estimated_resolution_angstrom,-0.05,fsc_threshold,'red','--') - plt.legend() - plt.tight_layout() - plt.show() - -def plot_plottile( - fsc_values: torch.Tensor, - frequency_pixels: torch.Tensor, - pixel_spacing_angstroms: float, - estimated_resolution_frequency_pixel: float, - fsc_threshold: float - ): - import plotille - - fig = plotille.Figure() - fig.width = 60 - fig.height = 20 - fig.set_x_limits(float(frequency_pixels[1]), float(frequency_pixels[-1])) - fig.set_y_limits(0, 1) - def resolution_callback(x,_): - return '{:.2f}'.format((1 / x) * pixel_spacing_angstroms) - fig.x_ticks_fkt = resolution_callback - fig.plot(frequency_pixels[1:].numpy(), fsc_values[1:].numpy(),lc='blue',label='FSC') - - fig.plot([float(frequency_pixels[1].numpy()),estimated_resolution_frequency_pixel],[fsc_threshold,fsc_threshold],lc='red',label=' ') - fig.plot([estimated_resolution_frequency_pixel,estimated_resolution_frequency_pixel],[0,fsc_threshold],lc='red',label=' ') - print(fig.show(legend=True)) \ No newline at end of file