smash some functions together

charon_vna/vna.py

@@ -305,7 +305,7 @@ class Charon:
 
         return np.mean(data[1] / data[0])
 
-    def sweep_b_over_a(self, **kwargs):
+    def capture(self, callback: Callable[int, int] | None = None):
         s = xr.DataArray(
             np.zeros(
                 [len(self.frequency), len(self.ports), len(self.ports)],
@@ -318,38 +318,35 @@ class Charon:
                 n=list(self.ports),
             ),
         )
+
+        total_count = len(s.m) * len(s.n) * len(s.frequency)
+        count = 0
+
         for m in s.m.data:
             for n in s.n.data:
                 self.set_switches(b=m - 1, a=n - 1)
-                s.loc[dict(m=m, n=n)] = self.vna_capture(frequency=self.frequency, **kwargs)
-        return s
 
-    def vna_capture(self, frequency: npt.ArrayLike, callback: Callable[int, int] | None = None):
-        s = xr.DataArray(
-            np.empty(len(frequency), dtype=np.complex128),
-            dims=["frequency"],
-            coords=dict(
-                frequency=frequency,
-            ),
-        )
-        for ff, freq in enumerate(s.frequency.data):
-            if callback is not None:
-                # report progress during sweep
-                callback(ff, len(s.frequency))
-            self.set_output(frequency=freq, power=-5)
-            self.sdr.rx_destroy_buffer()
-            self.sdr.rx_lo = int(freq)
-            self.sdr.rx_enabled_channels = [0, 1]
-            self.sdr.gain_control_mode_chan0 = "manual"
-            self.sdr.gain_control_mode_chan1 = "manual"
-            self.sdr.rx_hardwaregain_chan0 = 40
-            self.sdr.rx_hardwaregain_chan1 = 40
-            rx = self.sdr.rx()
-            s.loc[dict(frequency=freq)] = np.mean(rx[1] / rx[0])
+                for ff, freq in enumerate(s.frequency.data):
+                    if callback is not None:
+                        # report progress during sweep
+                        callback(count, total_count)
+
+                    self.set_output(frequency=freq, power=-5)
+                    self.sdr.rx_destroy_buffer()
+                    self.sdr.rx_lo = int(freq)
+                    self.sdr.rx_enabled_channels = [0, 1]
+                    self.sdr.gain_control_mode_chan0 = "manual"
+                    self.sdr.gain_control_mode_chan1 = "manual"
+                    self.sdr.rx_hardwaregain_chan0 = 40
+                    self.sdr.rx_hardwaregain_chan1 = 40
+                    rx = self.sdr.rx()
+                    s.loc[dict(frequency=freq, m=m, n=n)] = np.mean(rx[1] / rx[0])
+
+                    count += 1
 
         if callback is not None:
             # mark capture as complete
-            callback(len(s.frequency), len(s.frequency))
+            callback(total_count, total_count)
 
         return s
 

charon_vna/vna_dev.py

@@ -11,7 +11,7 @@ frequency = np.linspace(80e6, 280e6, 31)
 vna = Charon(frequency=frequency, ports=2)
 
 # %%
-s = vna.sweep_b_over_a()
+s = vna.capture()
 
 # %%
 for m in s.m.data: