From 0f0bc6b88fe8cb80b947e576cdfa4a7124982f54 Mon Sep 17 00:00:00 2001
From: ozpv <39195175+ozpv@users.noreply.github.com>
Date: Tue, 19 May 2026 01:48:04 -0500
Subject: add complement parameter

---
 src/retain.rs | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

(limited to 'src/retain.rs')

diff --git a/src/retain.rs b/src/retain.rs
index 8c4c8eb..5c14246 100644
--- a/src/retain.rs
+++ b/src/retain.rs
@@ -2,12 +2,26 @@ use num_complex::Complex;
 use rustc_hash::FxHashSet;
 
 #[inline(always)]
-pub fn retain_top_n_magnitudes(spectrum: &mut [Complex<f32>], n: usize) {
+pub fn retain_top_n_magnitudes(spectrum: &mut [Complex<f32>], n: usize, complement: bool) {
     // you'd be keeping the entire signal
-    // or keeping nothing at all
-    if n >= spectrum.len() {
+    if n >= spectrum.len() && !complement {
         return;
-    } else if n == 0 {
+    }
+
+    if n == 0 && complement {
+        return;
+    }
+
+    // keeping nothing at all
+    if n == 0 && !complement {
+        for phasor in spectrum {
+            *phasor = Complex::ZERO;
+        }
+
+        return;
+    }
+
+    if n >= spectrum.len() && complement {
         for phasor in spectrum {
             *phasor = Complex::ZERO;
         }
@@ -30,7 +44,11 @@ pub fn retain_top_n_magnitudes(spectrum: &mut [Complex<f32>], n: usize) {
         .collect::<FxHashSet<usize>>();
 
     for (i, phasor) in spectrum.iter_mut().enumerate() {
-        if !top_indices.contains(&i) {
+        if top_indices.contains(&i) && complement {
+            *phasor = Complex::ZERO;
+        }
+
+        if !top_indices.contains(&i) && !complement {
             *phasor = Complex::ZERO;
         }
     }
-- 
cgit v1.2.3