Refactor mcts checking

Author: hombit

src/data/multi_color_time_series.rs

@@ -1,12 +1,69 @@
 use crate::data::TimeSeries;
 use crate::float_trait::Float;
 use crate::multicolor::PassbandTrait;
+use crate::PassbandSet;
 
+use itertools::Either;
+use itertools::EitherOrBoth;
+use itertools::Itertools;
 use std::collections::BTreeMap;
 use std::ops::{Deref, DerefMut};
 
 pub struct MultiColorTimeSeries<'a, P: PassbandTrait, T: Float>(BTreeMap<P, TimeSeries<'a, T>>);
 
+impl<'a, P, T> MultiColorTimeSeries<'a, P, T>
+where
+    P: PassbandTrait,
+    T: Float,
+{
+    pub fn new(map: impl Into<BTreeMap<P, TimeSeries<'a, T>>>) -> Self {
+        Self(map.into())
+    }
+
+    pub fn iter_passband_set<'slf, 'ps>(
+        &'slf self,
+        passband_set: &'ps PassbandSet<P>,
+    ) -> impl Iterator<Item = (&P, Option<&TimeSeries<'a, T>>)> + 'ps
+    where
+        'a: 'ps,
+        'slf: 'ps,
+        'ps: 'slf,
+    {
+        match passband_set {
+            PassbandSet::AllAvailable => Either::Left(self.0.iter().map(|(p, ts)| (p, Some(ts)))),
+            PassbandSet::FixedSet(set) => Either::Right(set.iter().map(|p| (p, self.0.get(p)))),
+        }
+    }
+
+    pub fn iter_passband_set_mut<'slf, 'ps>(
+        &'slf mut self,
+        passband_set: &'ps PassbandSet<P>,
+    ) -> impl Iterator<Item = (&P, Option<&mut TimeSeries<'a, T>>)> + 'ps
+    where
+        'a: 'ps,
+        'slf: 'ps,
+        'ps: 'slf,
+    {
+        match passband_set {
+            PassbandSet::AllAvailable => {
+                Either::Left(self.0.iter_mut().map(|(p, ts)| (p, Some(ts))))
+            }
+            PassbandSet::FixedSet(set) => Either::Right(
+                set.iter()
+                    .merge_join_by(self.0.iter_mut(), |p1, (p2, _ts)| p1.cmp(p2))
+                    .filter_map(|either_or_both| match either_or_both {
+                        // mcts misses required passband
+                        EitherOrBoth::Left(p) => Some((p, None)),
+                        // mcts has some passban passband_set doesn't require
+                        EitherOrBoth::Right(_) => None,
+                        // passbands match
+                        EitherOrBoth::Both(p, (_, ts)) => Some((p, Some(ts))),
+                    }),
+            ),
+        }
+    }
+}
+
 impl<'a, P: PassbandTrait, T: Float> Deref for MultiColorTimeSeries<'a, P, T> {
     type Target = BTreeMap<P, TimeSeries<'a, T>>;
 

src/evaluator.rs

@@ -75,6 +75,42 @@ pub trait EvaluatorInfoTrait {
     fn is_variability_required(&self) -> bool {
         self.get_info().variability_required
     }
+
+    fn check_ts<F>(&self, ts: &mut TimeSeries<F>) -> Result<(), EvaluatorError>
+    where
+        F: Float,
+    {
+        self.check_ts_length(ts)?;
+        self.check_ts_variability(ts)
+    }
+
+    /// Checks if [TimeSeries] has enough points to evaluate the feature
+    fn check_ts_length<F>(&self, ts: &TimeSeries<F>) -> Result<(), EvaluatorError>
+    where
+        F: Float,
+    {
+        let length = ts.lenu();
+        if length < self.min_ts_length() {
+            Err(EvaluatorError::ShortTimeSeries {
+                actual: length,
+                minimum: self.min_ts_length(),
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Checks if [TimeSeries] meets variability requirement
+    fn check_ts_variability<F>(&self, ts: &mut TimeSeries<F>) -> Result<(), EvaluatorError>
+    where
+        F: Float,
+    {
+        if self.is_variability_required() && ts.is_plateau() {
+            Err(EvaluatorError::FlatTimeSeries)
+        } else {
+            Ok(())
+        }
+    }
 }
 
 // impl<P> EvaluatorInfoTrait for P
@@ -146,33 +182,6 @@ pub trait FeatureEvaluator<T: Float>:
             Err(_) => vec![fill_value; self.size_hint()],
         }
     }
-
-    fn check_ts(&self, ts: &mut TimeSeries<T>) -> Result<(), EvaluatorError> {
-        self.check_ts_length(ts)?;
-        self.check_ts_variability(ts)
-    }
-
-    /// Checks if [TimeSeries] has enough points to evaluate the feature
-    fn check_ts_length(&self, ts: &TimeSeries<T>) -> Result<(), EvaluatorError> {
-        let length = ts.lenu();
-        if length < self.min_ts_length() {
-            Err(EvaluatorError::ShortTimeSeries {
-                actual: length,
-                minimum: self.min_ts_length(),
-            })
-        } else {
-            Ok(())
-        }
-    }
-
-    /// Checks if [TimeSeries] meets variability requirement
-    fn check_ts_variability(&self, ts: &mut TimeSeries<T>) -> Result<(), EvaluatorError> {
-        if self.is_variability_required() && ts.is_plateau() {
-            Err(EvaluatorError::FlatTimeSeries)
-        } else {
-            Ok(())
-        }
-    }
 }
 
 pub trait OwnedArrays<T>

src/lib.rs

@@ -1,5 +1,7 @@
 #![doc = include_str!("../README.md")]
 
+extern crate core;
+
 #[cfg(test)]
 #[macro_use]
 mod tests;

src/multicolor/features/color_of_median.rs

@@ -52,7 +52,7 @@ where
 }
 
 lazy_info!(
-    COLOR_MEDIAN_INFO,
+    COLOR_OF_MEDIAN_INFO,
     size: 1,
     min_ts_length: 1,
     t_required: false,
@@ -67,7 +67,7 @@ where
     P: Ord,
 {
     fn get_info(&self) -> &EvaluatorInfo {
-        &COLOR_MEDIAN_INFO
+        &COLOR_OF_MEDIAN_INFO
     }
 }
 
@@ -98,11 +98,10 @@ where
     P: PassbandTrait,
     T: Float,
 {
-    fn eval_multicolor(
+    fn eval_multicolor_no_mcts_check(
         &self,
         mcts: &mut MultiColorTimeSeries<P, T>,
     ) -> Result<Vec<T>, MultiColorEvaluatorError> {
-        self.check_mcts_passabands(mcts)?;
         let mut medians = [T::zero(); 2];
         for (median, passband) in medians.iter_mut().zip(self.passbands.iter()) {
             *median = self

src/multicolor/monochrome_feature.rs

@@ -106,11 +106,10 @@ where
     T: Float,
     F: FeatureEvaluator<T>,
 {
-    fn eval_multicolor(
+    fn eval_multicolor_no_mcts_check(
         &self,
         mcts: &mut MultiColorTimeSeries<P, T>,
     ) -> Result<Vec<T>, MultiColorEvaluatorError> {
-        self.check_mcts_passabands(mcts)?;
         match &self.passband_set {
             PassbandSet::FixedSet(set) => set
                 .iter()