Flatter MultiColorTimeSeries

hombit · hombit · commit c440a5d2d782 · 2023-03-15T12:27:22.000-05:00
diff --git a/src/data/multi_color_time_series.rs b/src/data/multi_color_time_series.rs
@@ -1,52 +1,71 @@
 use crate::data::TimeSeries;
 use crate::float_trait::Float;
 use crate::multicolor::PassbandTrait;
-use crate::PassbandSet;
+use crate::{DataSample, 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>>);
+pub struct MultiColorTimeSeries<'a, P: PassbandTrait, T: Float> {
+    mapping: BTreeMap<P, TimeSeries<'a, T>>,
+    flat: Option<FlatMultiColorTimeSeries<'static, P, T>>,
+}
 
 impl<'a, 'p, P, T> MultiColorTimeSeries<'a, P, T>
 where
     P: PassbandTrait + 'p,
     T: Float,
 {
     pub fn new(map: impl Into<BTreeMap<P, TimeSeries<'a, T>>>) -> Self {
-        Self(map.into())
+        Self {
+            mapping: map.into(),
+            flat: None,
+        }
+    }
+
+    pub fn flatten(&mut self) -> &FlatMultiColorTimeSeries<'static, P, T> {
+        self.flat
+            .get_or_insert_with(|| FlatMultiColorTimeSeries::from_mapping(&mut self.mapping))
+    }
+
+    pub fn passbands<'slf>(
+        &'slf self,
+    ) -> std::collections::btree_map::Keys<'slf, P, TimeSeries<'a, T>>
+    where
+        'a: 'slf,
+    {
+        self.mapping.keys()
     }
 
     pub fn iter_passband_set<'slf, 'ps>(
         &'slf self,
         passband_set: &'ps PassbandSet<P>,
-    ) -> impl Iterator<Item = (&P, Option<&TimeSeries<'a, T>>)> + 'ps
+    ) -> impl Iterator<Item = (&P, Option<&TimeSeries<'a, T>>)> + 'slf
     where
-        'a: 'ps,
-        'slf: 'ps,
-        'ps: 'slf,
+        'a: 'slf,
+        'ps: 'a,
     {
         match passband_set {
-            PassbandSet::AllAvailable => Either::Left(self.0.iter().map(|(p, ts)| (p, Some(ts)))),
+            PassbandSet::AllAvailable => {
+                Either::Left(self.mapping.iter().map(|(p, ts)| (p, Some(ts))))
+            }
             PassbandSet::FixedSet(set) => Either::Right(self.iter_matched_passbands(set.iter())),
         }
     }
 
     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
+    ) -> impl Iterator<Item = (&P, Option<&mut TimeSeries<'a, T>>)> + 'slf
     where
-        'a: 'ps,
-        'slf: 'ps,
-        'ps: 'slf,
+        'a: 'slf,
+        'ps: 'a,
     {
         match passband_set {
             PassbandSet::AllAvailable => {
-                Either::Left(self.0.iter_mut().map(|(p, ts)| (p, Some(ts))))
+                Either::Left(self.mapping.iter_mut().map(|(p, ts)| (p, Some(ts))))
             }
             PassbandSet::FixedSet(set) => {
                 Either::Right(self.iter_matched_passbands_mut(set.iter()))
@@ -58,15 +77,15 @@ where
         &self,
         passband_it: impl Iterator<Item = &'p P>,
     ) -> impl Iterator<Item = (&'p P, Option<&TimeSeries<'a, T>>)> {
-        passband_it.map(|p| (p, self.0.get(p)))
+        passband_it.map(|p| (p, self.mapping.get(p)))
     }
 
     pub fn iter_matched_passbands_mut(
         &mut self,
         passband_it: impl Iterator<Item = &'p P>,
     ) -> impl Iterator<Item = (&'p P, Option<&mut TimeSeries<'a, T>>)> {
         passband_it
-            .merge_join_by(self.0.iter_mut(), |p1, (p2, _ts)| p1.cmp(p2))
+            .merge_join_by(self.mapping.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)),
@@ -78,24 +97,48 @@ where
     }
 }
 
-impl<'a, P: PassbandTrait, T: Float> Deref for MultiColorTimeSeries<'a, P, T> {
-    type Target = BTreeMap<P, TimeSeries<'a, T>>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
+impl<'a, P: PassbandTrait, T: Float> FromIterator<(P, TimeSeries<'a, T>)>
+    for MultiColorTimeSeries<'a, P, T>
+{
+    fn from_iter<I: IntoIterator<Item = (P, TimeSeries<'a, T>)>>(iter: I) -> Self {
+        Self {
+            mapping: iter.into_iter().collect(),
+            flat: None,
+        }
     }
 }
 
-impl<'a, P: PassbandTrait, T: Float> DerefMut for MultiColorTimeSeries<'a, P, T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
+pub struct FlatMultiColorTimeSeries<'a, P: PassbandTrait, T: Float> {
+    pub t: DataSample<'a, T>,
+    pub m: DataSample<'a, T>,
+    pub w: DataSample<'a, T>,
+    pub passbands: Vec<P>,
 }
 
-impl<'a, P: PassbandTrait, T: Float> FromIterator<(P, TimeSeries<'a, T>)>
-    for MultiColorTimeSeries<'a, P, T>
+impl<P, T> FlatMultiColorTimeSeries<'static, P, T>
+where
+    P: PassbandTrait,
+    T: Float,
 {
-    fn from_iter<I: IntoIterator<Item = (P, TimeSeries<'a, T>)>>(iter: I) -> Self {
-        Self(iter.into_iter().collect())
+    pub fn from_mapping(mapping: &mut BTreeMap<P, TimeSeries<T>>) -> Self {
+        let (t, m, w, passbands): (Vec<_>, Vec<_>, Vec<_>, _) = mapping
+            .iter_mut()
+            .map(|(p, ts)| {
+                itertools::multizip((
+                    ts.t.as_slice().iter().copied(),
+                    ts.m.as_slice().iter().copied(),
+                    ts.w.as_slice().iter().copied(),
+                    std::iter::repeat(p.clone()),
+                ))
+            })
+            .kmerge_by(|(t1, _, _, _), (t2, _, _, _)| t1 <= t2)
+            .multiunzip();
+
+        Self {
+            t: t.into(),
+            m: m.into(),
+            w: w.into(),
+            passbands,
+        }
     }
 }
diff --git a/src/multicolor/features/color_of_maximum.rs b/src/multicolor/features/color_of_maximum.rs
@@ -89,10 +89,14 @@ where
     P: PassbandTrait,
     T: Float,
 {
-    fn eval_multicolor_no_mcts_check(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+    {
         let mut maxima = [T::zero(); 2];
         for ((_passband, mcts), maximum) in mcts
             .iter_matched_passbands_mut(self.passbands.iter())
diff --git a/src/multicolor/features/color_of_median.rs b/src/multicolor/features/color_of_median.rs
@@ -98,10 +98,14 @@ where
     P: PassbandTrait,
     T: Float,
 {
-    fn eval_multicolor_no_mcts_check(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+    {
         let mut medians = [T::zero(); 2];
         for ((passband, mcts), median) in mcts
             .iter_matched_passbands_mut(self.passbands.iter())
diff --git a/src/multicolor/features/color_of_minimum.rs b/src/multicolor/features/color_of_minimum.rs
@@ -89,10 +89,14 @@ where
     P: PassbandTrait,
     T: Float,
 {
-    fn eval_multicolor_no_mcts_check(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+    {
         let mut minima = [T::zero(); 2];
         for ((_passband, mcts), minimum) in mcts
             .iter_matched_passbands_mut(self.passbands.iter())
diff --git a/src/multicolor/monochrome_feature.rs b/src/multicolor/monochrome_feature.rs
@@ -106,23 +106,26 @@ where
     T: Float,
     F: FeatureEvaluator<T>,
 {
-    fn eval_multicolor_no_mcts_check(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+    {
         match &self.passband_set {
-            PassbandSet::FixedSet(set) => set
-                .iter()
-                .map(|passband| {
-                    self.feature.eval(mcts.get_mut(passband).expect(
-                        "we checked all needed passbands are in mcts, but we still cannot find one",
-                    )).map_err(|error| MultiColorEvaluatorError::MonochromeEvaluatorError {
-                        passband: passband.name().into(),
-                        error,
-                    })
-                })
-                .flatten_ok()
-                .collect(),
+            PassbandSet::FixedSet(set) => {
+                mcts.iter_matched_passbands_mut(set.iter())
+                    .map(|(passband, ts)| {
+                        self.feature.eval_no_ts_check(
+                                ts.expect("we checked all needed passbands are in mcts, but we still cannot find one")
+                        ).map_err(|error| MultiColorEvaluatorError::MonochromeEvaluatorError {
+                            passband: passband.name().into(),
+                            error,
+                        })
+                    }).flatten_ok().collect()
+            }
             PassbandSet::AllAvailable => panic!("passband_set must be FixedSet variant here"),
         }
     }
diff --git a/src/multicolor/multicolor_evaluator.rs b/src/multicolor/multicolor_evaluator.rs
@@ -50,20 +50,22 @@ enum InternalMctsError {
 }
 
 impl InternalMctsError {
-    fn into_multi_color_evaluator_error<P, T>(
+    fn into_multi_color_evaluator_error<'mcts, 'a, 'ps, P, T>(
         self,
-        mcts: &MultiColorTimeSeries<P, T>,
-        ps: &PassbandSet<P>,
+        mcts: &'mcts MultiColorTimeSeries<'a, P, T>,
+        ps: &'ps PassbandSet<P>,
     ) -> MultiColorEvaluatorError
     where
+        'ps: 'a,
+        'a: 'mcts,
         P: PassbandTrait,
         T: Float,
     {
         match self {
             InternalMctsError::MultiColorEvaluatorError(e) => e,
             InternalMctsError::InternalWrongPassbandSet => {
                 MultiColorEvaluatorError::wrong_passbands_error(
-                    mcts.keys(),
+                    mcts.passbands(),
                     match ps {
                         PassbandSet::FixedSet(ps) => ps.iter(),
                         PassbandSet::AllAvailable => {
@@ -88,37 +90,55 @@ where
     T: Float,
 {
     /// Version of [MultiColorEvaluator::eval_multicolor] without basic [MultiColorTimeSeries] checks
-    fn eval_multicolor_no_mcts_check(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError>;
+    fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts;
 
     /// Vector of feature values or `EvaluatorError`
-    fn eval_multicolor(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    fn eval_multicolor<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+        P: 'a,
+    {
         self.check_mcts(mcts)?;
         self.eval_multicolor_no_mcts_check(mcts)
     }
 
     /// Returns vector of feature values and fill invalid components with given value
-    fn eval_or_fill_multicolor(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
+    fn eval_or_fill_multicolor<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
         fill_value: T,
-    ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+    ) -> Result<Vec<T>, MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+        P: 'a,
+    {
         Ok(match self.eval_multicolor(mcts) {
             Ok(v) => v,
             Err(_) => vec![fill_value; self.size_hint()],
         })
     }
 
     /// Check [MultiColorTimeSeries] to have required passbands and individual [TimeSeries] are valid
-    fn check_mcts(
-        &self,
-        mcts: &mut MultiColorTimeSeries<P, T>,
-    ) -> Result<(), MultiColorEvaluatorError> {
+    fn check_mcts<'slf, 'a, 'mcts>(
+        &'slf self,
+        mcts: &'mcts mut MultiColorTimeSeries<'a, P, T>,
+    ) -> Result<(), MultiColorEvaluatorError>
+    where
+        'slf: 'a,
+        'a: 'mcts,
+        P: 'a,
+    {
         mcts.iter_passband_set_mut(self.get_passband_set())
             .map(|(p, maybe_ts)| {
                 maybe_ts
@@ -187,10 +207,14 @@ mod tests {
     where
         T: Float,
     {
-        fn eval_multicolor_no_mcts_check(
-            &self,
-            _mcts: &mut MultiColorTimeSeries<MonochromePassband<'static, f64>, T>,
-        ) -> Result<Vec<T>, MultiColorEvaluatorError> {
+        fn eval_multicolor_no_mcts_check<'slf, 'a, 'mcts>(
+            &'slf self,
+            _mcts: &'mcts mut MultiColorTimeSeries<'a, MonochromePassband<'static, f64>, T>,
+        ) -> Result<Vec<T>, MultiColorEvaluatorError>
+        where
+            'slf: 'a,
+            'a: 'mcts,
+        {
             Ok(vec![T::zero()])
         }
     }
diff --git a/src/multicolor/multicolor_extractor.rs b/src/multicolor/multicolor_extractor.rs
