new naming conventions

mjo22 · Dec 18, 2024 · d8b74b1 · d8b74b1
1 parent 1bf6fa4
commit d8b74b1
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/src/cryojax/simulator/_transfer_theory/base_transfer_theory.py b/src/cryojax/simulator/_transfer_theory/base_transfer_theory.py
@@ -18,7 +18,6 @@ def __call__(
         frequency_grid_in_angstroms: Float[Array, "y_dim x_dim 2"],
         *,
         voltage_in_kilovolts: Float[Array, ""] | float = 300.0,
-        defocus_offset: Float[Array, ""] | float = 0.0,
     ) -> Float[Array, "y_dim x_dim"] | Complex[Array, "y_dim x_dim"]:
         raise NotImplementedError
 
@@ -29,7 +28,7 @@ class AbstractTransferTheory(Module, strict=True):
     @abstractmethod
     def __call__(
         self,
-        fourier_phase_or_wavefunction_at_exit_plane: (
+        array: (
             Complex[
                 Array,
                 "{instrument_config.padded_y_dim} "
@@ -41,7 +40,6 @@ def __call__(
             ]
         ),
         instrument_config: InstrumentConfig,
-        defocus_offset: Float[Array, ""] | float = 0.0,
     ) -> (
         Complex[
             Array,
@@ -51,5 +49,5 @@ def __call__(
             Array, "{instrument_config.padded_y_dim} {instrument_config.padded_x_dim}"
         ]
     ):
-        """Pass an image through the transfer theory."""
+        """Pass a quantity through the transfer theory."""
         raise NotImplementedError
diff --git a/src/cryojax/simulator/_transfer_theory/contrast_transfer_theory.py b/src/cryojax/simulator/_transfer_theory/contrast_transfer_theory.py
@@ -128,7 +128,7 @@ def __init__(
     @override
     def __call__(
         self,
-        fourier_phase_at_exit_plane: Complex[
+        object_spectrum_at_exit_plane: Complex[
             Array,
             "{instrument_config.padded_y_dim} {instrument_config.padded_x_dim//2+1}",
         ],
@@ -145,6 +145,6 @@ def __call__(
         )
         # ... compute the contrast as the CTF multiplied by the exit plane
         # phase shifts
-        fourier_contrast_at_detector_plane = ctf_array * fourier_phase_at_exit_plane
+        contrast_spectrum_at_detector_plane = ctf_array * object_spectrum_at_exit_plane
 
-        return fourier_contrast_at_detector_plane
+        return contrast_spectrum_at_detector_plane
diff --git a/src/cryojax/simulator/_transfer_theory/wave_transfer_theory.py b/src/cryojax/simulator/_transfer_theory/wave_transfer_theory.py
@@ -115,7 +115,7 @@ def __init__(
     @override
     def __call__(
         self,
-        fourier_wavefunction_at_exit_plane: Complex[
+        wavefunction_spectrum_at_exit_plane: Complex[
             Array,
             "{instrument_config.padded_y_dim} {instrument_config.padded_x_dim}",
         ],
@@ -132,8 +132,8 @@ def __call__(
         )
         # ... compute the contrast as the CTF multiplied by the exit plane
         # phase shifts
-        fourier_wavefunction_at_detector_plane = (
-            wtf_array * fourier_wavefunction_at_exit_plane
+        wavefunction_spectrum_at_detector_plane = (
+            wtf_array * wavefunction_spectrum_at_exit_plane
         )
 
-        return fourier_wavefunction_at_detector_plane
+        return wavefunction_spectrum_at_detector_plane