Merge pull request #63 from jdickerson95/jd_b-env

Calculate b-factor envelope and CTF separately, and update `torch-fourier-filter` dependency accordingly

Author: mgiammar

pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
     "roma",
     "tqdm",
     "torch-fourier-slice>=v0.2.0",
-    "torch-fourier-filter>=v0.2.3",
+    "torch-fourier-filter>=v0.2.6",
     "torch-so3>=v0.2.0",
     "ttsim3d>=v0.4.0",
     "lmfit",

src/leopard_em/pydantic_models/utils.py

@@ -5,6 +5,7 @@
 
 import torch
 from torch_fourier_filter.ctf import calculate_ctf_2d
+from torch_fourier_filter.envelopes import b_envelope
 
 # Using the TYPE_CHECKING statement to avoid circular imports
 if TYPE_CHECKING:
@@ -166,13 +167,22 @@ def calculate_ctf_filter_stack_full_args(
             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,
         )
+        # calc B-envelope and apply
+        b_envelope_tmp = b_envelope(
+            B=kwargs["ctf_B_factor"],
+            image_shape=template_shape,
+            pixel_size=kwargs["pixel_size"],
+            rfft=True,
+            fftshift=False,
+            device=tmp.device,
+        )
+        tmp *= b_envelope_tmp
         ctf_list.append(tmp)
 
     ctf = torch.stack(ctf_list, dim=0)