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 num_modes: usize,
26 alpha: f64,
28 #[serde(default = "defaults::init")]
30 init: Init,
31 #[serde(default = "defaults::tolerance")]
33 tolerance: f64,
34 #[serde(default = "defaults::tau")]
36 tau: f64,
37 #[serde(default = "defaults::max_iterations")]
39 max_iterations: usize,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
42 noise_channels: Option<usize>,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
45 noise_sd: Option<f64>,
46 #[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#[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 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
135fn 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
162pub 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);