scirs2_fft/
backend.rs

1//! FFT Backend System
2//!
3//! This module provides a pluggable backend system for FFT implementations,
4//! similar to SciPy's backend model. This allows users to choose between
5//! different FFT implementations at runtime.
6
7use crate::error::{FFTError, FFTResult};
8use rustfft::FftPlanner;
9use scirs2_core::numeric::Complex64;
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock};
12
13/// FFT Backend trait for implementing different FFT algorithms
14pub trait FftBackend: Send + Sync {
15    /// Name of the backend
16    fn name(&self) -> &str;
17
18    /// Description of the backend
19    fn description(&self) -> &str;
20
21    /// Check if this backend is available
22    fn is_available(&self) -> bool;
23
24    /// Perform forward FFT
25    fn fft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()>;
26
27    /// Perform inverse FFT
28    fn ifft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()>;
29
30    /// Perform FFT with specific size (may be cached)
31    fn fft_sized(
32        &self,
33        input: &[Complex64],
34        output: &mut [Complex64],
35        size: usize,
36    ) -> FFTResult<()>;
37
38    /// Perform inverse FFT with specific size (may be cached)
39    fn ifft_sized(
40        &self,
41        input: &[Complex64],
42        output: &mut [Complex64],
43        size: usize,
44    ) -> FFTResult<()>;
45
46    /// Check if this backend supports a specific feature
47    fn supports_feature(&self, feature: &str) -> bool;
48}
49
50/// RustFFT backend implementation
51pub struct RustFftBackend {
52    planner: Arc<Mutex<FftPlanner<f64>>>,
53}
54
55impl RustFftBackend {
56    pub fn new() -> Self {
57        Self {
58            planner: Arc::new(Mutex::new(FftPlanner::new())),
59        }
60    }
61}
62
63impl Default for RustFftBackend {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl FftBackend for RustFftBackend {
70    fn name(&self) -> &str {
71        "rustfft"
72    }
73
74    fn description(&self) -> &str {
75        "Pure Rust FFT implementation using RustFFT library"
76    }
77
78    fn is_available(&self) -> bool {
79        true
80    }
81
82    fn fft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()> {
83        self.fft_sized(input, output, input.len())
84    }
85
86    fn ifft(&self, input: &[Complex64], output: &mut [Complex64]) -> FFTResult<()> {
87        self.ifft_sized(input, output, input.len())
88    }
89
90    fn fft_sized(
91        &self,
92        input: &[Complex64],
93        output: &mut [Complex64],
94        size: usize,
95    ) -> FFTResult<()> {
96        if input.len() != size || output.len() != size {
97            return Err(FFTError::ValueError(
98                "Input and output sizes must match the specified size".to_string(),
99            ));
100        }
101
102        // Get cached plan from the planner
103        let mut planner = self.planner.lock().unwrap();
104        let fft = planner.plan_fft_forward(size);
105
106        // Convert to rustfft's Complex type
107        let mut buffer: Vec<rustfft::num_complex::Complex<f64>> = input
108            .iter()
109            .map(|&c| rustfft::num_complex::Complex::new(c.re, c.im))
110            .collect();
111
112        // Perform FFT
113        fft.process(&mut buffer);
114
115        // Copy to output
116        for (i, &c) in buffer.iter().enumerate() {
117            output[i] = Complex64::new(c.re, c.im);
118        }
119
120        Ok(())
121    }
122
123    fn ifft_sized(
124        &self,
125        input: &[Complex64],
126        output: &mut [Complex64],
127        size: usize,
128    ) -> FFTResult<()> {
129        if input.len() != size || output.len() != size {
130            return Err(FFTError::ValueError(
131                "Input and output sizes must match the specified size".to_string(),
132            ));
133        }
134
135        // Get cached plan from the planner
136        let mut planner = self.planner.lock().unwrap();
137        let fft = planner.plan_fft_inverse(size);
138
139        // Convert to rustfft's Complex type
140        let mut buffer: Vec<rustfft::num_complex::Complex<f64>> = input
141            .iter()
142            .map(|&c| rustfft::num_complex::Complex::new(c.re, c.im))
143            .collect();
144
145        // Perform IFFT
146        fft.process(&mut buffer);
147
148        // Copy to output with normalization
149        let scale = 1.0 / size as f64;
150        for (i, &c) in buffer.iter().enumerate() {
151            output[i] = Complex64::new(c.re * scale, c.im * scale);
152        }
153
154        Ok(())
155    }
156
157    fn supports_feature(&self, feature: &str) -> bool {
158        matches!(feature, "1d_fft" | "2d_fft" | "nd_fft" | "cached_plans")
159    }
160}
161
162/// Backend manager for FFT operations
163pub struct BackendManager {
164    backends: Arc<Mutex<HashMap<String, Arc<dyn FftBackend>>>>,
165    current_backend: Arc<Mutex<String>>,
166}
167
168impl BackendManager {
169    /// Create a new backend manager with default backends
170    pub fn new() -> Self {
171        let mut backends = HashMap::new();
172
173        // Add default RustFFT backend
174        let rustfft_backend = Arc::new(RustFftBackend::new()) as Arc<dyn FftBackend>;
175        backends.insert("rustfft".to_string(), rustfft_backend);
176
177        Self {
178            backends: Arc::new(Mutex::new(backends)),
179            current_backend: Arc::new(Mutex::new("rustfft".to_string())),
180        }
181    }
182
183    /// Register a new backend
184    pub fn register_backend(&self, name: String, backend: Arc<dyn FftBackend>) -> FFTResult<()> {
185        let mut backends = self.backends.lock().unwrap();
186        if backends.contains_key(&name) {
187            return Err(FFTError::ValueError(format!(
188                "Backend '{name}' already exists"
189            )));
190        }
191        backends.insert(name, backend);
192        Ok(())
193    }
194
195    /// Get available backends
196    pub fn list_backends(&self) -> Vec<String> {
197        let backends = self.backends.lock().unwrap();
198        backends.keys().cloned().collect()
199    }
200
201    /// Set the current backend
202    pub fn set_backend(&self, name: &str) -> FFTResult<()> {
203        let backends = self.backends.lock().unwrap();
204        if !backends.contains_key(name) {
205            return Err(FFTError::ValueError(format!("Backend '{name}' not found")));
206        }
207
208        // Check if backend is available
209        if let Some(backend) = backends.get(name) {
210            if !backend.is_available() {
211                return Err(FFTError::ValueError(format!(
212                    "Backend '{name}' is not available"
213                )));
214            }
215        }
216
217        *self.current_backend.lock().unwrap() = name.to_string();
218        Ok(())
219    }
220
221    /// Get current backend name
222    pub fn get_backend_name(&self) -> String {
223        self.current_backend.lock().unwrap().clone()
224    }
225
226    /// Get current backend
227    pub fn get_backend(&self) -> Arc<dyn FftBackend> {
228        let current_name = self.current_backend.lock().unwrap();
229        let backends = self.backends.lock().unwrap();
230        backends
231            .get(&*current_name)
232            .cloned()
233            .expect("Current backend should always exist")
234    }
235
236    /// Get backend info
237    pub fn get_backend_info(&self, name: &str) -> Option<BackendInfo> {
238        let backends = self.backends.lock().unwrap();
239        backends.get(name).map(|backend| BackendInfo {
240            name: backend.name().to_string(),
241            description: backend.description().to_string(),
242            available: backend.is_available(),
243        })
244    }
245}
246
247impl Default for BackendManager {
248    fn default() -> Self {
249        Self::new()
250    }
251}
252
253/// Information about a backend
254#[derive(Debug, Clone)]
255pub struct BackendInfo {
256    pub name: String,
257    pub description: String,
258    pub available: bool,
259}
260
261impl std::fmt::Display for BackendInfo {
262    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263        write!(
264            f,
265            "{} - {} ({})",
266            self.name,
267            self.description,
268            if self.available {
269                "available"
270            } else {
271                "not available"
272            }
273        )
274    }
275}
276
277/// Global backend manager instance
278static GLOBAL_BACKEND_MANAGER: OnceLock<BackendManager> = OnceLock::new();
279
280/// Get the global backend manager
281#[allow(dead_code)]
282pub fn get_backend_manager() -> &'static BackendManager {
283    GLOBAL_BACKEND_MANAGER.get_or_init(BackendManager::new)
284}
285
286/// Initialize global backend manager with custom configuration
287#[allow(dead_code)]
288pub fn init_backend_manager(manager: BackendManager) -> Result<(), &'static str> {
289    GLOBAL_BACKEND_MANAGER
290        .set(manager)
291        .map_err(|_| "Global backend _manager already initialized")
292}
293
294/// List available backends
295#[allow(dead_code)]
296pub fn list_backends() -> Vec<String> {
297    get_backend_manager().list_backends()
298}
299
300/// Set the current backend
301#[allow(dead_code)]
302pub fn set_backend(name: &str) -> FFTResult<()> {
303    get_backend_manager().set_backend(name)
304}
305
306/// Get current backend name
307#[allow(dead_code)]
308pub fn get_backend_name() -> String {
309    get_backend_manager().get_backend_name()
310}
311
312/// Get backend information
313#[allow(dead_code)]
314pub fn get_backend_info(name: &str) -> Option<BackendInfo> {
315    get_backend_manager().get_backend_info(name)
316}
317
318/// Context manager for temporarily using a different backend
319pub struct BackendContext {
320    previous_backend: String,
321    manager: &'static BackendManager,
322}
323
324impl BackendContext {
325    /// Create a new backend context
326    pub fn new(_backendname: &str) -> FFTResult<Self> {
327        let manager = get_backend_manager();
328        let previous_backend = manager.get_backend_name();
329
330        // Set the new backend
331        manager.set_backend(_backendname)?;
332
333        Ok(Self {
334            previous_backend,
335            manager,
336        })
337    }
338}
339
340impl Drop for BackendContext {
341    fn drop(&mut self) {
342        // Restore previous backend
343        let _ = self.manager.set_backend(&self.previous_backend);
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350
351    #[test]
352    fn test_rustfft_backend() {
353        let backend = RustFftBackend::new();
354        assert_eq!(backend.name(), "rustfft");
355        assert!(backend.is_available());
356        assert!(backend.supports_feature("1d_fft"));
357    }
358
359    #[test]
360    fn test_backend_manager() {
361        let manager = BackendManager::new();
362
363        // Check default backend
364        assert_eq!(manager.get_backend_name(), "rustfft");
365
366        // List backends
367        let backends = manager.list_backends();
368        assert!(backends.contains(&"rustfft".to_string()));
369
370        // Get backend info
371        let info = manager.get_backend_info("rustfft").unwrap();
372        assert!(info.available);
373    }
374
375    #[test]
376    fn test_backend_context() {
377        let manager = get_backend_manager();
378        let original = manager.get_backend_name();
379
380        {
381            let _ctx = BackendContext::new("rustfft").unwrap();
382            assert_eq!(manager.get_backend_name(), "rustfft");
383        }
384
385        // Backend should be restored
386        assert_eq!(manager.get_backend_name(), original);
387    }
388}