From d8b74b1e82a1229ec45cd751a12ea9f50d4dcffc Mon Sep 17 00:00:00 2001
From: Michael O'Brien <michaelobrien@g.harvard.edu>
Date: Wed, 18 Dec 2024 18:04:00 -0500
Subject: [PATCH] new naming conventions

---
 .../simulator/_transfer_theory/base_transfer_theory.py    | 6 ++----
 .../_transfer_theory/contrast_transfer_theory.py          | 6 +++---
 .../simulator/_transfer_theory/wave_transfer_theory.py    | 8 ++++----
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/cryojax/simulator/_transfer_theory/base_transfer_theory.py b/src/cryojax/simulator/_transfer_theory/base_transfer_theory.py
index 1a944514..1469f063 100644
--- 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
index acbd2ffa..f334adc8 100644
--- 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
index 12ae1296..3eabd1bd 100644
--- 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