Merge pull request #263 from mjo22/defocus-convention

Change defocus parameter convention to mean defocus
mjo22 · Sep 8, 2024 · dfbea8a · dfbea8a
2 parents 35a666f + 497a475
commit dfbea8a
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/src/cryojax/simulator/_transfer_theory/contrast_transfer_theory.py b/src/cryojax/simulator/_transfer_theory/contrast_transfer_theory.py
@@ -5,7 +5,7 @@
 from equinox import field
 from jaxtyping import Array, Complex, Float
 
-from ..._errors import error_if_negative, error_if_not_fractional, error_if_not_positive
+from ..._errors import error_if_negative, error_if_not_fractional
 from ...constants import convert_keV_to_angstroms
 from ...image.operators import (
     Constant,
@@ -19,6 +19,21 @@
 class ContrastTransferFunction(AbstractTransferFunction, strict=True):
     """Compute an astigmatic Contrast Transfer Function (CTF) with a
     spherical aberration correction and amplitude contrast ratio.
+
+    !!! info
+        `cryojax` uses a convention different from CTFFIND for
+        astigmatism parameters. It returns defocus major and minor
+        axes, called "defocus1" and "defocus2". In order to convert
+        from CTFFIND to `cryojax`,
+
+        ```python
+        defocus1, defocus2 = ... # Read from CTFFIND
+        ctf = ContrastTransferFunction(
+            defocus_in_angstroms=(defocus1+defocus2)/2,
+            astigmatism_in_angstroms=defocus1-defocus2,
+            ...
+        )
+        ```
     """
 
     defocus_in_angstroms: Float[Array, ""]
@@ -41,8 +56,8 @@ def __init__(
     ):
         """**Arguments:**
 
-        - `defocus_u_in_angstroms`: The major axis defocus in Angstroms.
-        - `defocus_v_in_angstroms`: The minor axis defocus in Angstroms.
+        - `defocus_in_angstroms`: The mean defocus in Angstroms.
+        - `astigmatism_in_angstroms`: The amount of astigmatism in Angstroms.
         - `astigmatism_angle`: The defocus angle.
         - `voltage_in_kilovolts`:
             The accelerating voltage in kV. This field is treated as *static*, i.e.
@@ -53,7 +68,7 @@ def __init__(
         - `amplitude_contrast_ratio`: The amplitude contrast ratio.
         - `phase_shift`: The additional phase shift.
         """
-        self.defocus_in_angstroms = error_if_not_positive(defocus_in_angstroms)
+        self.defocus_in_angstroms = error_if_negative(defocus_in_angstroms)
         self.astigmatism_in_angstroms = jnp.asarray(astigmatism_in_angstroms)
         self.astigmatism_angle = jnp.asarray(astigmatism_angle)
         self.voltage_in_kilovolts = voltage_in_kilovolts
@@ -81,13 +96,15 @@ def __call__(
             )
         else:
             wavelength_in_angstroms = jnp.asarray(wavelength_in_angstroms)
-        defocus_axis_1_in_angstroms = self.defocus_in_angstroms + jnp.asarray(
-            defocus_offset
+        defocus_axis_1_in_angstroms = (
+            self.defocus_in_angstroms
+            + jnp.asarray(defocus_offset)
+            + self.astigmatism_in_angstroms / 2
         )
         defocus_axis_2_in_angstroms = (
             self.defocus_in_angstroms
-            + self.astigmatism_in_angstroms
             + jnp.asarray(defocus_offset)
+            - self.astigmatism_in_angstroms / 2
         )
         # Compute phase shifts for CTF
         phase_shifts = compute_phase_shifts_with_amplitude_contrast_ratio(

diff --git a/tests/test_agree_with_cistem.py b/tests/test_agree_with_cistem.py
@@ -50,8 +50,8 @@ def test_ctf_with_cistem(defocus1, defocus2, asti_angle, kV, cs, ac, pixel_size)
         k_sqr, theta = cartesian_to_polar(freqs, square=True)
         # Compute cryojax CTF
         optics = ContrastTransferFunction(
-            defocus_in_angstroms=defocus1,
-            astigmatism_in_angstroms=defocus2 - defocus1,
+            defocus_in_angstroms=(defocus1 + defocus2) / 2,
+            astigmatism_in_angstroms=defocus1 - defocus2,
             astigmatism_angle=asti_angle,
             voltage_in_kilovolts=kV,
             spherical_aberration_in_mm=cs,