From ac6b10bed6b60a1d053bf09d517ad08535c5398d Mon Sep 17 00:00:00 2001
From: Daniel Kohler <11864045+ddkohler@users.noreply.github.com>
Date: Thu, 19 Sep 2024 12:10:01 -0500
Subject: [PATCH] Update heterodyne.py

---
 tests/mixed/heterodyne.py | 39 +++++++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/tests/mixed/heterodyne.py b/tests/mixed/heterodyne.py
index 4610b60..829684a 100644
--- a/tests/mixed/heterodyne.py
+++ b/tests/mixed/heterodyne.py
@@ -42,10 +42,6 @@
 
 def check_phase():
     # TODO: use assertions to check phase
-    ...
-    
-
-if __name__ == "__main__":
     scan = exp.run(ham, mp=False, windowed=True)
 
     efields = scan.efields()
@@ -55,26 +51,37 @@ def check_phase():
 
     driven_sig = 1j * efields.prod(axis=-2)
 
-    # plot amplitude    
-    import matplotlib.pyplot as plt
-    plt.close('all')
-    fig, gs = wt.artists.create_figure(width='single', nrows=2, cols=[1, 1, 1])
+    diff = []
+    driven_diff = []
+
     for i in range(scan.npulses):
         lo = efields[..., i, :]
         if scan.pm[i] == 1:
             lo = lo.conjugate()
-        diff = (lo * sig).imag.sum(axis=-1)
-        driven_diff = (lo * driven_sig).imag.sum(axis=-1)
+        diff.append((lo * sig).imag.sum(axis=-1))
+        driven_diff.append((lo * driven_sig).imag.sum(axis=-1))
+
+    return diff, driven_diff
+
+
+if __name__ == "__main__":
+    diff, driven_diff = check_phase()
+
+    # plot amplitude    
+    import matplotlib.pyplot as plt
+    plt.close('all')
+    fig, gs = wt.artists.create_figure(width='single', nrows=2, cols=[1, 1, 1])
+    for i, (di, ddi) in enumerate(zip(diff, driven_diff)):
 
         axi = plt.subplot(gs[i])
-        axi.pcolormesh(exp.d2.points, exp.d1.points, -diff, 
-                    vmin=-np.abs(diff).max(),
-                    vmax=np.abs(diff).max(),
+        axi.pcolormesh(exp.d2.points, exp.d1.points, -di, 
+                    vmin=-np.abs(di).max(),
+                    vmax=np.abs(di).max(),
                     cmap='signed')
         axi2 = plt.subplot(gs[i+3])
-        axi2.pcolormesh(exp.d2.points, exp.d1.points, driven_diff,
-                        vmin=-np.abs(driven_diff).max(),
-                        vmax=np.abs(driven_diff).max(),
+        axi2.pcolormesh(exp.d2.points, exp.d1.points, ddi,
+                        vmin=-np.abs(ddi).max(),
+                        vmax=np.abs(ddi).max(),
                         cmap='signed')
         [ax.grid(True) for ax in [axi, axi2]]