complex calib add

tests/test_processing.py

@@ -27,10 +27,13 @@ from rfg_adc_plotter.gui.pyqtgraph_backend import (
 )
 from rfg_adc_plotter.processing.calibration import (
     build_calib_envelope,
+    build_complex_calibration_curve,
     calibrate_freqs,
     load_calib_envelope,
+    load_complex_calibration,
     recalculate_calibration_c,
     save_calib_envelope,
+    save_complex_calibration,
 )
 from rfg_adc_plotter.processing.background import (
     load_fft_background,
@@ -49,7 +52,9 @@ from rfg_adc_plotter.processing.fft import (
 )
 from rfg_adc_plotter.processing.normalization import (
     build_calib_envelopes,
+    fit_complex_calibration_to_width,
     normalize_by_calib,
+    normalize_by_complex_calibration,
     normalize_by_envelope,
     resample_envelope,
 )
@@ -148,6 +153,46 @@ class ProcessingTests(unittest.TestCase):
             with self.assertRaises(ValueError):
                 load_calib_envelope(path)
 
+    def test_complex_calibration_curve_roundtrip(self):
+        ch1 = np.asarray([1.0, 2.0, 3.0], dtype=np.float32)
+        ch2 = np.asarray([0.5, -1.0, 4.0], dtype=np.float32)
+        curve = build_complex_calibration_curve(ch1, ch2)
+        expected = np.asarray([1.0 + 0.5j, 2.0 - 1.0j, 3.0 + 4.0j], dtype=np.complex64)
+
+        self.assertTrue(np.iscomplexobj(curve))
+        self.assertTrue(np.allclose(curve, expected))
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = os.path.join(tmp_dir, "complex_calibration")
+            saved_path = save_complex_calibration(path, curve)
+            loaded = load_complex_calibration(saved_path)
+            self.assertTrue(saved_path.endswith(".npy"))
+            self.assertEqual(loaded.dtype, np.complex64)
+            self.assertTrue(np.allclose(loaded, expected))
+
+    def test_fit_complex_calibration_to_width_pads_or_trims(self):
+        calib = np.asarray([1.0 + 1.0j, 2.0 + 2.0j], dtype=np.complex64)
+        padded = fit_complex_calibration_to_width(calib, 4)
+        trimmed = fit_complex_calibration_to_width(
+            np.asarray([1.0 + 1.0j, 2.0 + 2.0j, 3.0 + 3.0j], dtype=np.complex64),
+            2,
+        )
+
+        self.assertEqual(padded.shape, (4,))
+        self.assertTrue(np.allclose(padded, np.asarray([1.0 + 1.0j, 2.0 + 2.0j, 1.0 + 0.0j, 1.0 + 0.0j], dtype=np.complex64)))
+        self.assertEqual(trimmed.shape, (2,))
+        self.assertTrue(np.allclose(trimmed, np.asarray([1.0 + 1.0j, 2.0 + 2.0j], dtype=np.complex64)))
+
+    def test_normalize_by_complex_calibration_handles_zero_and_length_mismatch(self):
+        signal = np.asarray([2.0 + 2.0j, 4.0 + 0.0j, 3.0 + 3.0j], dtype=np.complex64)
+        calib = np.asarray([1.0 + 1.0j, 0.0 + 0.0j], dtype=np.complex64)
+        normalized = normalize_by_complex_calibration(signal, calib)
+        expected = np.asarray([2.0 + 0.0j, 4.0 + 0.0j, 3.0 + 3.0j], dtype=np.complex64)
+
+        self.assertTrue(np.iscomplexobj(normalized))
+        self.assertTrue(np.all(np.isfinite(normalized)))
+        self.assertTrue(np.allclose(normalized, expected))
+
     def test_fft_background_roundtrip_and_rejects_non_1d_payload(self):
         background = np.asarray([0.5, 1.5, 2.5], dtype=np.float32)
         with tempfile.TemporaryDirectory() as tmp_dir: