complex: richer API

dsp/src/complex.rs

@@ -1,17 +1,58 @@
-use super::atan2;
+use super::{atan2, cossin};
 use serde::{Deserialize, Serialize};
 
-#[derive(Copy, Clone, Default, Deserialize, Serialize)]
+#[derive(Copy, Clone, Default, PartialEq, Debug, Deserialize, Serialize)]
 pub struct Complex<T>(pub T, pub T);
 
 impl Complex<i32> {
-    pub fn power(&self) -> i32 {
-        (((self.0 as i64) * (self.0 as i64)
-            + (self.1 as i64) * (self.1 as i64))
-            >> 32) as i32
+    /// Return a Complex on the unit circle given an angle.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// use dsp::Complex;
+    /// Complex::<i32>::from_angle(0);
+    /// Complex::<i32>::from_angle(1 << 30);  // pi/2
+    /// Complex::<i32>::from_angle(-1 << 30);  // -pi/2
+    /// ```
+    #[inline(always)]
+    pub fn from_angle(angle: i32) -> Complex<i32> {
+        cossin(angle)
     }
 
-    pub fn phase(&self) -> i32 {
+    /// Return the absolute square (the squared magnitude).
+    ///
+    /// Note: Normalization is `1 << 31`, i.e. Q0.31.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// use dsp::Complex;
+    /// assert_eq!(Complex(i32::MAX, 0).abs_sqr(), i32::MAX - 1);
+    /// assert_eq!(Complex(i32::MIN + 1, 0).abs_sqr(), i32::MAX - 1);
+    /// ```
+    pub fn abs_sqr(&self) -> i32 {
+        (((self.0 as i64) * (self.0 as i64)
+            + (self.1 as i64) * (self.1 as i64))
+            >> 31) as i32
+    }
+
+    /// Return the angle.
+    ///
+    /// Note: Normalization is `1 << 31 == pi`.
+    ///
+    /// Example:
+    ///
+    /// ```
+    /// use dsp::Complex;
+    /// assert_eq!(Complex(i32::MAX, 0).arg(), 0);
+    /// assert_eq!(Complex(-i32::MAX, 1).arg(), i32::MAX);
+    /// assert_eq!(Complex(-i32::MAX, -1).arg(), -i32::MAX);
+    /// assert_eq!(Complex(0, -i32::MAX).arg(), -i32::MAX >> 1);
+    /// assert_eq!(Complex(0, i32::MAX).arg(), (i32::MAX >> 1) + 1);
+    /// assert_eq!(Complex(i32::MAX, i32::MAX).arg(), (i32::MAX >> 2) + 1);
+    /// ```
+    pub fn arg(&self) -> i32 {
         atan2(self.1, self.0)
     }
 }

dsp/src/lockin.rs

@@ -1,4 +1,4 @@
-use super::{cossin, iir_int, Complex};
+use super::{iir_int, Complex};
 use serde::{Deserialize, Serialize};
 
 #[derive(Copy, Clone, Default, Deserialize, Serialize)]
@@ -19,7 +19,7 @@ impl Lockin {
 
     pub fn update(&mut self, signal: i32, phase: i32) -> Complex<i32> {
         // Get the LO signal for demodulation.
-        let m = cossin(phase);
+        let m = Complex::from_angle(phase);
 
         // Mix with the LO signal, filter with the IIR lowpass,
         // return IQ (in-phase and quadrature) data.
@@ -35,6 +35,23 @@ impl Lockin {
             ),
         )
     }
+
+    pub fn feed<I: IntoIterator<Item = i32>>(
+        &mut self,
+        signal: I,
+        phase: i32,
+        frequency: i32,
+    ) -> Option<Complex<i32>> {
+        let mut phase = phase;
+
+        signal
+            .into_iter()
+            .map(|s| {
+                phase = phase.wrapping_add(frequency);
+                self.update(s, phase)
+            })
+            .last()
+    }
 }
 
 #[cfg(test)]

src/bin/lockin.rs

@@ -123,27 +123,26 @@ const APP: () = {
             22, // relative PLL phase bandwidth: 2**-22, TODO: expose
         );
 
-        // Harmonic index of the LO: -1 to _de_modulate the fundamental
+        // Harmonic index of the LO: -1 to _de_modulate the fundamental (complex conjugate)
         let harmonic: i32 = -1;
         // Demodulation LO phase offset
         let phase_offset: i32 = 0;
         let sample_frequency = ((pll_frequency >> SAMPLE_BUFFER_SIZE_LOG2)
             as i32)
             .wrapping_mul(harmonic); // TODO: maybe rounding bias
-        let mut sample_phase =
+        let sample_phase =
             phase_offset.wrapping_add(pll_phase.wrapping_mul(harmonic));
 
-        for i in 0..adc_samples[0].len() {
-            // Convert to signed, MSB align the ADC sample.
-            let input = (adc_samples[0][i] as i16 as i32) << 16;
-            // Obtain demodulated, filtered IQ sample.
-            let output = lockin.update(input, sample_phase);
-            // Advance the sample phase.
-            sample_phase = sample_phase.wrapping_add(sample_frequency);
-
+        if let Some(output) = lockin.feed(
+            adc_samples[0].iter().map(|&i|
+                // Convert to signed, MSB align the ADC sample.
+                (i as i16 as i32) << 16),
+            sample_phase,
+            sample_frequency,
+        ) {
             // Convert from IQ to power and phase.
-            let mut power = output.power() as _;
-            let mut phase = output.phase() as _;
+            let mut power = output.abs_sqr() as _;
+            let mut phase = output.arg() as _;
 
             // Filter power and phase through IIR filters.
             // Note: Normalization to be done in filters. Phase will wrap happily.
@@ -152,13 +151,17 @@ const APP: () = {
                 phase = iir_ch[1][j].update(&mut iir_state[1][j], phase);
             }
 
+            // TODO: IIR filter DC gain needs to be 1/(1 << 16)
+
             // Note(unsafe): range clipping to i16 is ensured by IIR filters above.
             // Convert to DAC data.
-            unsafe {
-                dac_samples[0][i] =
-                    power.to_int_unchecked::<i16>() as u16 ^ 0x8000;
-                dac_samples[1][i] =
-                    phase.to_int_unchecked::<i16>() as u16 ^ 0x8000;
+            for i in 0..dac_samples[0].len() {
+                unsafe {
+                    dac_samples[0][i] =
+                        power.to_int_unchecked::<i16>() as u16 ^ 0x8000;
+                    dac_samples[1][i] =
+                        phase.to_int_unchecked::<i16>() as u16 ^ 0x8000;
+                }
             }
         }
     }