modewise_operations/
mvmd.rs

1use modewise_decomposition::ModeDecomposer;
2use modewise_mvmd::{FrequencyInit, Mvmd, MvmdParams, NaMvmd, NoiseParams, Vmd};
3use modewise_pipeline::{
4    Artifact, ArtifactKind, Inputs, Operation, OperationDescriptor, Outputs, Port, resolve_params,
5};
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10use crate::input;
11
12#[derive(Serialize, Deserialize, JsonSchema)]
13#[serde(rename_all = "lowercase")]
14enum Init {
15    Zero,
16    Linear,
17    Exponential,
18    Uniform,
19}
20
21#[derive(Serialize, Deserialize, JsonSchema)]
22#[serde(deny_unknown_fields)]
23struct Params {
24    /// Number of modes K.
25    num_modes: usize,
26    /// Bandwidth penalty; larger means narrower modes.
27    alpha: f64,
28    /// Starting centre frequencies: `zero`, `linear` (evenly over [0, 0.5], as Dmocrito/mvmd), `exponential`, or `uniform` ((0.5/K)·k, as the original MATLAB VMD and MVMD codes and vmdpy).
29    #[serde(default = "defaults::init")]
30    init: Init,
31    /// Stop when the mean squared update falls below this.
32    #[serde(default = "defaults::tolerance")]
33    tolerance: f64,
34    /// Dual ascent step; 0 lets noise stay out of the modes.
35    #[serde(default = "defaults::tau")]
36    tau: f64,
37    /// Iteration cap.
38    #[serde(default = "defaults::max_iterations")]
39    max_iterations: usize,
40    /// na_mvmd only: white-noise channels appended before decomposing.
41    #[serde(default, skip_serializing_if = "Option::is_none")]
42    noise_channels: Option<usize>,
43    /// na_mvmd only: noise SD relative to the unit-SD signal channels.
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    noise_sd: Option<f64>,
46    /// na_mvmd only: seed for the noise.
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    seed: Option<u64>,
49}
50
51mod defaults {
52    use modewise_mvmd::MvmdParams;
53
54    fn reference() -> MvmdParams {
55        MvmdParams::new(1, 1.0)
56    }
57    pub fn init() -> super::Init {
58        super::Init::Zero
59    }
60    pub fn tolerance() -> f64 {
61        reference().tolerance
62    }
63    pub fn tau() -> f64 {
64        reference().tau
65    }
66    pub fn max_iterations() -> usize {
67        reference().max_iterations
68    }
69}
70
71/// The methods that share the MVMD solver and its parameters.
72#[derive(Clone, Copy, PartialEq)]
73enum Variant {
74    Mvmd,
75    Vmd,
76    NaMvmd,
77}
78
79impl Params {
80    fn mvmd_params(&self) -> MvmdParams {
81        let init = match self.init {
82            Init::Zero => FrequencyInit::Zero,
83            Init::Linear => FrequencyInit::Linear,
84            Init::Exponential => FrequencyInit::Exponential,
85            Init::Uniform => FrequencyInit::Uniform,
86        };
87        MvmdParams {
88            num_modes: self.num_modes,
89            alpha: self.alpha,
90            init,
91            tolerance: self.tolerance,
92            tau: self.tau,
93            max_iterations: self.max_iterations,
94        }
95    }
96
97    /// Fills the noise defaults for na_mvmd, and rejects noise settings elsewhere.
98    fn complete(mut self, variant: Variant) -> Result<Self, String> {
99        let has_noise =
100            self.noise_channels.is_some() || self.noise_sd.is_some() || self.seed.is_some();
101        match variant {
102            Variant::NaMvmd => {
103                let d = NoiseParams::default();
104                self.noise_channels.get_or_insert(d.channels);
105                self.noise_sd.get_or_insert(d.sd);
106                self.seed.get_or_insert(d.seed);
107            }
108            _ if has_noise => {
109                return Err("`noise_channels`, `noise_sd` and `seed` only apply to na_mvmd".into());
110            }
111            _ => {}
112        }
113        Ok(self)
114    }
115
116    fn method(&self, variant: Variant) -> Result<Box<dyn ModeDecomposer + Send + Sync>, String> {
117        let params = self.mvmd_params();
118        let method: Box<dyn ModeDecomposer + Send + Sync> = match variant {
119            Variant::Mvmd => Box::new(Mvmd::new(params).map_err(|e| e.to_string())?),
120            Variant::Vmd => Box::new(Vmd::new(params).map_err(|e| e.to_string())?),
121            Variant::NaMvmd => {
122                let d = NoiseParams::default();
123                let noise = NoiseParams {
124                    channels: self.noise_channels.unwrap_or(d.channels),
125                    sd: self.noise_sd.unwrap_or(d.sd),
126                    seed: self.seed.unwrap_or(d.seed),
127                };
128                Box::new(NaMvmd::new(params, noise).map_err(|e| e.to_string())?)
129            }
130        };
131        Ok(method)
132    }
133}
134
135/// The shared schema narrowed to what `variant` accepts, with the noise defaults filled in by
136/// the same completion that `resolve` uses.
137fn variant_schema(variant: Variant) -> Value {
138    let mut schema = modewise_pipeline::schema_of::<Params>();
139    let noisy = matches!(variant, Variant::NaMvmd);
140    let defaults =
141        serde_json::from_value::<Params>(serde_json::json!({"num_modes": 1, "alpha": 1.0}))
142            .ok()
143            .and_then(|p| p.complete(variant).ok())
144            .and_then(|p| serde_json::to_value(p).ok())
145            .unwrap_or_default();
146    if let Some(properties) = schema.get_mut("properties").and_then(Value::as_object_mut) {
147        properties.retain(|field, _| {
148            noisy || !matches!(field.as_str(), "noise_channels" | "noise_sd" | "seed")
149        });
150        for field in ["noise_channels", "noise_sd", "seed"] {
151            if let (Some(default), Some(spec)) = (
152                defaults.get(field),
153                properties.get_mut(field).and_then(Value::as_object_mut),
154            ) {
155                spec.insert("default".into(), default.clone());
156            }
157        }
158    }
159    schema
160}
161
162/// `mvmd`, `vmd` (each channel on its own) and `na_mvmd` (noise-assisted).
163pub struct MvmdOperation;
164pub struct VmdOperation;
165pub struct NaMvmdOperation;
166
167fn descriptor(name: &'static str) -> OperationDescriptor {
168    OperationDescriptor {
169        name,
170        version: env!("CARGO_PKG_VERSION"),
171        inputs: &[Port {
172            name: "signal",
173            kinds: &[ArtifactKind::Series],
174        }],
175        outputs: &[Port {
176            name: "modes",
177            kinds: &[ArtifactKind::Decomposition],
178        }],
179    }
180}
181
182fn resolve(variant: Variant, params: Value) -> Result<Value, String> {
183    let (typed, _) = resolve_params::<Params>(params)?;
184    let typed = typed.complete(variant)?;
185    typed.method(variant)?;
186    serde_json::to_value(&typed).map_err(|e| e.to_string())
187}
188
189fn run(variant: Variant, params: &Value, inputs: &Inputs) -> Result<Outputs, String> {
190    let method = serde_json::from_value::<Params>(params.clone())
191        .map_err(|e| e.to_string())?
192        .method(variant)?;
193    let Artifact::Series(signal) = input(inputs, "signal")? else {
194        return Err("input \"signal\" must be a series".into());
195    };
196    let modes = method.decompose(signal).map_err(|e| e.to_string())?;
197    Ok(Outputs::from([(
198        "modes".into(),
199        Artifact::Decomposition(modes),
200    )]))
201}
202
203macro_rules! mvmd_family {
204    ($operation:ident, $name:literal, $variant:expr, $summary:literal) => {
205        impl Operation for $operation {
206            fn summary(&self) -> &'static str {
207                $summary
208            }
209
210            fn params_schema(&self) -> serde_json::Value {
211                variant_schema($variant)
212            }
213
214            fn descriptor(&self) -> OperationDescriptor {
215                descriptor($name)
216            }
217
218            fn resolve(&self, params: Value) -> Result<Value, String> {
219                resolve($variant, params)
220            }
221
222            fn run(&self, params: &Value, inputs: &Inputs) -> Result<Outputs, String> {
223                run($variant, params, inputs)
224            }
225        }
226    };
227}
228
229mvmd_family!(
230    MvmdOperation,
231    "mvmd",
232    Variant::Mvmd,
233    "Multivariate variational mode decomposition: K modes shared by all channels (Rehman and Aftab 2019)."
234);
235mvmd_family!(
236    VmdOperation,
237    "vmd",
238    Variant::Vmd,
239    "Variational mode decomposition of each channel on its own (Dragomiretskiy and Zosso 2014)."
240);
241mvmd_family!(
242    NaMvmdOperation,
243    "na_mvmd",
244    Variant::NaMvmd,
245    "Noise-assisted MVMD: MVMD with white-noise channels appended, less sensitive to K (Zisou et al. 2021)."
246);