1use scirs2_core::ndarray::{Array2, ArrayView2};
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9use sklears_core::{
10 error::{Result as SklResult, SklearsError},
11 traits::{Estimator, Fit, Transform},
12 types::Float,
13};
14use std::collections::HashMap;
15use std::fmt::Debug;
16use std::sync::{Arc, RwLock};
17
18static PLUGIN_REGISTRY: once_cell::sync::Lazy<RwLock<PluginRegistry>> =
20 once_cell::sync::Lazy::new(|| RwLock::new(PluginRegistry::new()));
21
22pub trait ManifoldPlugin: Send + Sync + Debug {
24 fn name(&self) -> &str;
26
27 fn version(&self) -> &str;
29
30 fn description(&self) -> &str;
32
33 fn author(&self) -> &str;
35
36 fn create_default(&self) -> Box<dyn CustomManifoldLearner>;
38
39 fn create_with_params(
41 &self,
42 params: &PluginParameters,
43 ) -> SklResult<Box<dyn CustomManifoldLearner>>;
44
45 fn default_parameters(&self) -> PluginParameters;
47
48 fn validate_parameters(&self, params: &PluginParameters) -> SklResult<()>;
50
51 fn metadata(&self) -> PluginMetadata {
53 PluginMetadata {
55 name: self.name().to_string(),
56 version: self.version().to_string(),
57 description: self.description().to_string(),
58 author: self.author().to_string(),
59 supported_features: self.supported_features(),
60 parameter_schema: self.parameter_schema(),
61 }
62 }
63
64 fn supported_features(&self) -> Vec<PluginFeature> {
66 vec![PluginFeature::DimensionalityReduction]
67 }
68
69 fn parameter_schema(&self) -> Vec<ParameterDefinition>;
71}
72
73pub trait CustomManifoldLearner: Send + Sync + Debug {
75 fn set_parameter(&mut self, name: &str, value: ParameterValue) -> SklResult<()>;
77
78 fn get_parameter(&self, name: &str) -> Option<ParameterValue>;
80
81 fn get_all_parameters(&self) -> HashMap<String, ParameterValue>;
83
84 fn fit(&mut self, x: &ArrayView2<Float>) -> SklResult<()>;
86
87 fn transform(&self, x: &ArrayView2<Float>) -> SklResult<Array2<Float>>;
89
90 fn fit_transform(&mut self, x: &ArrayView2<Float>) -> SklResult<Array2<Float>> {
92 self.fit(x)?;
93 self.transform(x)
94 }
95
96 fn is_fitted(&self) -> bool;
98
99 fn get_metadata(&self) -> CustomModelMetadata;
101
102 fn clone_learner(&self) -> Box<dyn CustomManifoldLearner>;
104}
105
106#[derive(Debug)]
108pub struct PluginRegistry {
109 plugins: HashMap<String, Arc<dyn ManifoldPlugin>>,
110}
111
112impl Default for PluginRegistry {
113 fn default() -> Self {
114 Self::new()
115 }
116}
117
118impl PluginRegistry {
119 pub fn new() -> Self {
121 Self {
122 plugins: HashMap::new(),
123 }
124 }
125
126 pub fn register_plugin(&mut self, plugin: Arc<dyn ManifoldPlugin>) -> SklResult<()> {
128 let name = plugin.name().to_string();
129
130 if self.plugins.contains_key(&name) {
131 return Err(SklearsError::InvalidInput(format!(
132 "Plugin '{}' is already registered",
133 name
134 )));
135 }
136
137 self.plugins.insert(name, plugin);
138 Ok(())
139 }
140
141 pub fn unregister_plugin(&mut self, name: &str) -> SklResult<()> {
143 if self.plugins.remove(name).is_none() {
144 return Err(SklearsError::InvalidInput(format!(
145 "Plugin '{}' is not registered",
146 name
147 )));
148 }
149 Ok(())
150 }
151
152 pub fn get_plugin(&self, name: &str) -> Option<Arc<dyn ManifoldPlugin>> {
154 self.plugins.get(name).cloned()
155 }
156
157 pub fn list_plugins(&self) -> Vec<String> {
159 self.plugins.keys().cloned().collect()
160 }
161
162 pub fn get_all_metadata(&self) -> Vec<PluginMetadata> {
164 self.plugins
165 .values()
166 .map(|plugin| plugin.metadata())
167 .collect()
168 }
169
170 pub fn create_instance(
172 &self,
173 name: &str,
174 params: Option<&PluginParameters>,
175 ) -> SklResult<Box<dyn CustomManifoldLearner>> {
176 let plugin = self
177 .get_plugin(name)
178 .ok_or_else(|| SklearsError::InvalidInput(format!("Plugin '{}' not found", name)))?;
179
180 match params {
181 Some(params) => plugin.create_with_params(params),
182 None => Ok(plugin.create_default()),
183 }
184 }
185}
186
187impl PluginRegistry {
189 pub fn global() -> &'static RwLock<PluginRegistry> {
191 &PLUGIN_REGISTRY
192 }
193}
194
195#[derive(Debug, Clone, PartialEq)]
197#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
198pub struct PluginParameters {
199 parameters: HashMap<String, ParameterValue>,
200}
201
202impl PluginParameters {
203 pub fn new() -> Self {
205 Self {
206 parameters: HashMap::new(),
207 }
208 }
209
210 pub fn set<T: Into<ParameterValue>>(&mut self, name: &str, value: T) -> &mut Self {
212 self.parameters.insert(name.to_string(), value.into());
213 self
214 }
215
216 pub fn get(&self, name: &str) -> Option<&ParameterValue> {
218 self.parameters.get(name)
219 }
220
221 pub fn contains(&self, name: &str) -> bool {
223 self.parameters.contains_key(name)
224 }
225
226 pub fn all(&self) -> &HashMap<String, ParameterValue> {
228 &self.parameters
229 }
230
231 pub fn merge(&mut self, other: &PluginParameters) {
233 for (key, value) in &other.parameters {
234 self.parameters.insert(key.clone(), value.clone());
235 }
236 }
237}
238
239impl Default for PluginParameters {
240 fn default() -> Self {
241 Self::new()
242 }
243}
244
245#[derive(Debug, Clone, PartialEq)]
247#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
248pub enum ParameterValue {
249 Int(i64),
251 Float(f64),
253 String(String),
255 Bool(bool),
257 IntArray(Vec<i64>),
259 FloatArray(Vec<f64>),
261 StringArray(Vec<String>),
263}
264
265impl From<i64> for ParameterValue {
266 fn from(value: i64) -> Self {
267 ParameterValue::Int(value)
268 }
269}
270
271impl From<i32> for ParameterValue {
272 fn from(value: i32) -> Self {
273 ParameterValue::Int(value as i64)
274 }
275}
276
277impl From<usize> for ParameterValue {
278 fn from(value: usize) -> Self {
279 ParameterValue::Int(value as i64)
280 }
281}
282
283impl From<f64> for ParameterValue {
284 fn from(value: f64) -> Self {
285 ParameterValue::Float(value)
286 }
287}
288
289impl From<f32> for ParameterValue {
290 fn from(value: f32) -> Self {
291 ParameterValue::Float(value as f64)
292 }
293}
294
295impl From<String> for ParameterValue {
296 fn from(value: String) -> Self {
297 ParameterValue::String(value)
298 }
299}
300
301impl From<&str> for ParameterValue {
302 fn from(value: &str) -> Self {
303 ParameterValue::String(value.to_string())
304 }
305}
306
307impl From<bool> for ParameterValue {
308 fn from(value: bool) -> Self {
309 ParameterValue::Bool(value)
310 }
311}
312
313impl From<Vec<i64>> for ParameterValue {
314 fn from(value: Vec<i64>) -> Self {
315 ParameterValue::IntArray(value)
316 }
317}
318
319impl From<Vec<f64>> for ParameterValue {
320 fn from(value: Vec<f64>) -> Self {
321 ParameterValue::FloatArray(value)
322 }
323}
324
325impl From<Vec<String>> for ParameterValue {
326 fn from(value: Vec<String>) -> Self {
327 ParameterValue::StringArray(value)
328 }
329}
330
331#[derive(Debug, Clone, PartialEq)]
333#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
334pub enum PluginFeature {
335 DimensionalityReduction,
337 Clustering,
339 Classification,
341 Regression,
343 Visualization,
345 OutOfSample,
347 IncrementalLearning,
349 Parallelization,
351 GPU,
353}
354
355#[derive(Debug, Clone, PartialEq)]
357#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
358pub struct ParameterDefinition {
359 pub name: String,
361 pub param_type: ParameterType,
363 pub description: String,
365 pub default_value: Option<ParameterValue>,
367 pub required: bool,
369 pub constraints: Option<ParameterConstraints>,
371}
372
373#[derive(Debug, Clone, PartialEq)]
375#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
376pub enum ParameterType {
377 Int,
379 Float,
381 String,
383 Bool,
385 IntArray,
387 FloatArray,
389 StringArray,
391}
392
393#[derive(Debug, Clone, PartialEq)]
395#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
396pub struct ParameterConstraints {
397 pub min_value: Option<f64>,
399 pub max_value: Option<f64>,
401 pub allowed_values: Option<Vec<String>>,
403 pub min_length: Option<usize>,
405 pub max_length: Option<usize>,
407}
408
409#[derive(Debug, Clone, PartialEq)]
411#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
412pub struct PluginMetadata {
413 pub name: String,
415 pub version: String,
417 pub description: String,
419 pub author: String,
421 pub supported_features: Vec<PluginFeature>,
423 pub parameter_schema: Vec<ParameterDefinition>,
425}
426
427#[derive(Debug, Clone, PartialEq)]
429#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
430pub struct CustomModelMetadata {
431 pub plugin_name: String,
433 pub plugin_version: String,
435 pub is_fitted: bool,
437 pub n_samples: Option<usize>,
439 pub n_features: Option<usize>,
441 pub n_components: Option<usize>,
443 pub training_time: Option<f64>,
445 pub parameters: HashMap<String, ParameterValue>,
447}
448
449#[derive(Debug)]
451pub struct CustomManifoldWrapper {
452 learner: Box<dyn CustomManifoldLearner>,
453 plugin_name: String,
454}
455
456impl CustomManifoldWrapper {
457 pub fn new(plugin_name: &str, params: Option<&PluginParameters>) -> SklResult<Self> {
459 let registry = PLUGIN_REGISTRY.read().expect("operation should succeed");
460 let learner = registry.create_instance(plugin_name, params)?;
461
462 Ok(Self {
463 learner,
464 plugin_name: plugin_name.to_string(),
465 })
466 }
467
468 pub fn learner(&self) -> &dyn CustomManifoldLearner {
470 self.learner.as_ref()
471 }
472
473 pub fn learner_mut(&mut self) -> &mut dyn CustomManifoldLearner {
475 self.learner.as_mut()
476 }
477
478 pub fn plugin_name(&self) -> &str {
480 &self.plugin_name
481 }
482}
483
484impl Clone for CustomManifoldWrapper {
485 fn clone(&self) -> Self {
486 Self {
487 learner: self.learner.clone_learner(),
488 plugin_name: self.plugin_name.clone(),
489 }
490 }
491}
492
493impl Estimator for CustomManifoldWrapper {
495 type Config = PluginParameters;
496 type Error = SklearsError;
497 type Float = Float;
498
499 fn config(&self) -> &Self::Config {
500 static EMPTY_CONFIG: once_cell::sync::Lazy<PluginParameters> =
502 once_cell::sync::Lazy::new(PluginParameters::new);
503 &EMPTY_CONFIG
504 }
505}
506
507impl Fit<ArrayView2<'_, Float>, ()> for CustomManifoldWrapper {
508 type Fitted = CustomManifoldWrapper;
509
510 fn fit(mut self, x: &ArrayView2<'_, Float>, _y: &()) -> SklResult<Self::Fitted> {
511 self.learner.fit(x)?;
512 Ok(self)
513 }
514}
515
516impl Transform<ArrayView2<'_, Float>, Array2<Float>> for CustomManifoldWrapper {
517 fn transform(&self, x: &ArrayView2<'_, Float>) -> SklResult<Array2<Float>> {
518 self.learner.transform(x)
519 }
520}
521
522pub mod utils {
524 use super::*;
525
526 pub fn register_plugin(plugin: Arc<dyn ManifoldPlugin>) -> SklResult<()> {
528 PLUGIN_REGISTRY
529 .write()
530 .expect("operation should succeed")
531 .register_plugin(plugin)
532 }
533
534 pub fn unregister_plugin(name: &str) -> SklResult<()> {
536 PLUGIN_REGISTRY
537 .write()
538 .expect("operation should succeed")
539 .unregister_plugin(name)
540 }
541
542 pub fn list_plugins() -> Vec<String> {
544 PLUGIN_REGISTRY
545 .read()
546 .expect("operation should succeed")
547 .list_plugins()
548 }
549
550 pub fn get_plugin_metadata(name: &str) -> Option<PluginMetadata> {
552 PLUGIN_REGISTRY
554 .read()
555 .expect("operation should succeed")
556 .get_plugin(name)
557 .map(|p| p.metadata())
558 }
559
560 pub fn get_all_plugin_metadata() -> Vec<PluginMetadata> {
562 PLUGIN_REGISTRY
563 .read()
564 .expect("operation should succeed")
565 .get_all_metadata()
566 }
567
568 pub fn create_plugin_instance(
570 name: &str,
571 params: Option<&PluginParameters>,
572 ) -> SklResult<CustomManifoldWrapper> {
573 CustomManifoldWrapper::new(name, params)
574 }
575
576 pub fn validate_parameters(plugin_name: &str, params: &PluginParameters) -> SklResult<()> {
578 let registry = PLUGIN_REGISTRY.read().expect("operation should succeed");
579 let plugin = registry.get_plugin(plugin_name).ok_or_else(|| {
580 SklearsError::InvalidInput(format!("Plugin '{}' not found", plugin_name))
581 })?;
582 plugin.validate_parameters(params)
583 }
584}
585
586#[allow(non_snake_case)]
587#[cfg(test)]
588mod tests {
589 use super::*;
590 use scirs2_core::ndarray::{Array2, ArrayView2};
591 use scirs2_core::random::thread_rng;
592
593 #[derive(Debug)]
595 struct ExamplePlugin;
596
597 impl ManifoldPlugin for ExamplePlugin {
598 fn name(&self) -> &str {
599 "ExamplePlugin"
600 }
601 fn version(&self) -> &str {
602 "1.0.0"
603 }
604 fn description(&self) -> &str {
605 "An example plugin for testing"
606 }
607 fn author(&self) -> &str {
608 "Test Author"
609 }
610
611 fn create_default(&self) -> Box<dyn CustomManifoldLearner> {
612 Box::new(ExampleLearner::default())
613 }
614
615 fn create_with_params(
616 &self,
617 params: &PluginParameters,
618 ) -> SklResult<Box<dyn CustomManifoldLearner>> {
619 let mut learner = ExampleLearner::default();
620
621 if let Some(ParameterValue::Int(n_components)) = params.get("n_components") {
622 learner.set_parameter("n_components", ParameterValue::Int(*n_components))?;
623 }
624
625 Ok(Box::new(learner))
626 }
627
628 fn default_parameters(&self) -> PluginParameters {
629 let mut params = PluginParameters::new();
630 params.set("n_components", 2i64);
631 params
632 }
633
634 fn validate_parameters(&self, params: &PluginParameters) -> SklResult<()> {
635 if let Some(ParameterValue::Int(n_components)) = params.get("n_components") {
636 if *n_components <= 0 {
637 return Err(SklearsError::InvalidInput(
638 "n_components must be positive".to_string(),
639 ));
640 }
641 }
642 Ok(())
643 }
644
645 fn parameter_schema(&self) -> Vec<ParameterDefinition> {
646 vec![ParameterDefinition {
647 name: "n_components".to_string(),
648 param_type: ParameterType::Int,
649 description: "Number of components".to_string(),
650 default_value: Some(ParameterValue::Int(2)),
651 required: false,
652 constraints: Some(ParameterConstraints {
653 min_value: Some(1.0),
654 max_value: None,
655 allowed_values: None,
656 min_length: None,
657 max_length: None,
658 }),
659 }]
660 }
661 }
662
663 #[derive(Debug, Clone)]
665 struct ExampleLearner {
666 n_components: usize,
667 fitted: bool,
668 embedding: Option<Array2<Float>>,
669 }
670
671 impl Default for ExampleLearner {
672 fn default() -> Self {
673 Self {
674 n_components: 2,
675 fitted: false,
676 embedding: None,
677 }
678 }
679 }
680
681 impl CustomManifoldLearner for ExampleLearner {
682 fn set_parameter(&mut self, name: &str, value: ParameterValue) -> SklResult<()> {
683 match name {
684 "n_components" => {
685 if let ParameterValue::Int(val) = value {
686 self.n_components = val as usize;
687 Ok(())
688 } else {
689 Err(SklearsError::InvalidInput(
690 "n_components must be an integer".to_string(),
691 ))
692 }
693 }
694 _ => Err(SklearsError::InvalidInput(format!(
695 "Unknown parameter: {}",
696 name
697 ))),
698 }
699 }
700
701 fn get_parameter(&self, name: &str) -> Option<ParameterValue> {
702 match name {
703 "n_components" => Some(ParameterValue::Int(self.n_components as i64)),
704 _ => None,
705 }
706 }
707
708 fn get_all_parameters(&self) -> HashMap<String, ParameterValue> {
709 let mut params = HashMap::new();
710 params.insert(
711 "n_components".to_string(),
712 ParameterValue::Int(self.n_components as i64),
713 );
714 params
715 }
716
717 fn fit(&mut self, x: &ArrayView2<Float>) -> SklResult<()> {
718 let (n_samples, _) = x.dim();
719
720 let mut rng = thread_rng();
722 let mut embedding = Array2::zeros((n_samples, self.n_components));
723 for elem in embedding.iter_mut() {
724 *elem = rng.random();
725 }
726
727 self.embedding = Some(embedding);
728 self.fitted = true;
729 Ok(())
730 }
731
732 fn transform(&self, _x: &ArrayView2<Float>) -> SklResult<Array2<Float>> {
733 if !self.fitted {
734 return Err(SklearsError::InvalidInput(
735 "Model is not fitted".to_string(),
736 ));
737 }
738
739 self.embedding
742 .clone()
743 .ok_or_else(|| SklearsError::InvalidInput("No embedding available".to_string()))
744 }
745
746 fn is_fitted(&self) -> bool {
747 self.fitted
748 }
749
750 fn get_metadata(&self) -> CustomModelMetadata {
751 CustomModelMetadata {
753 plugin_name: "ExamplePlugin".to_string(),
754 plugin_version: "1.0.0".to_string(),
755 is_fitted: self.fitted,
756 n_samples: self.embedding.as_ref().map(|e| e.nrows()),
757 n_features: None,
758 n_components: Some(self.n_components),
759 training_time: None,
760 parameters: self.get_all_parameters(),
761 }
762 }
763
764 fn clone_learner(&self) -> Box<dyn CustomManifoldLearner> {
765 Box::new(self.clone())
766 }
767 }
768
769 #[test]
770 fn test_plugin_registration() {
771 let plugin = Arc::new(ExamplePlugin);
772 let result = utils::register_plugin(plugin);
773 assert!(result.is_ok());
774
775 let plugins = utils::list_plugins();
776 assert!(plugins.contains(&"ExamplePlugin".to_string()));
777
778 utils::unregister_plugin("ExamplePlugin").expect("operation should succeed");
780 }
781
782 #[test]
783 fn test_plugin_instance_creation() {
784 let plugin = Arc::new(ExamplePlugin);
785 utils::register_plugin(plugin).expect("operation should succeed");
786
787 let wrapper = utils::create_plugin_instance("ExamplePlugin", None);
788 assert!(wrapper.is_ok());
789
790 utils::unregister_plugin("ExamplePlugin").expect("operation should succeed");
792 }
793
794 #[test]
795 fn test_parameter_validation() {
796 let plugin = Arc::new(ExamplePlugin);
797 utils::register_plugin(plugin).expect("operation should succeed");
798
799 let mut params = PluginParameters::new();
800 params.set("n_components", 5i64);
801
802 let result = utils::validate_parameters("ExamplePlugin", ¶ms);
803 assert!(result.is_ok());
804
805 params.set("n_components", -1i64);
807 let result = utils::validate_parameters("ExamplePlugin", ¶ms);
808 assert!(result.is_err());
809
810 utils::unregister_plugin("ExamplePlugin").expect("operation should succeed");
812 }
813}