use crate::classifier::{kNN, RandomForest};
use crate::control::Control;
use crate::gain::{ChangeInMean, ClassifierGain};
use crate::optimizer::{GridSearch, TwoStepSearch};
use crate::segmentation::{Segmentation, SegmentationType};
use crate::{BinarySegmentationResult, BinarySegmentationTree};
use ndarray;
pub fn changeforest(
X: &ndarray::ArrayView2<'_, f64>,
method: &str,
segmentation_type: &str,
control: &Control,
) -> BinarySegmentationResult {
let segmentation_type_enum: SegmentationType;
let mut tree: BinarySegmentationTree;
if segmentation_type == "bs" {
segmentation_type_enum = SegmentationType::BS;
} else if segmentation_type == "sbs" {
segmentation_type_enum = SegmentationType::SBS;
} else if segmentation_type == "wbs" {
segmentation_type_enum = SegmentationType::WBS;
} else {
panic!("segmentation_type must be one of 'bs', 'sbs', 'wbs'")
}
if method == "knn" {
let classifier = kNN::new(X, control);
let gain = ClassifierGain { classifier };
let optimizer = TwoStepSearch { gain };
let mut segmentation = Segmentation::new(segmentation_type_enum, &optimizer);
tree = BinarySegmentationTree::new(X);
tree.grow(&mut segmentation);
BinarySegmentationResult::from_tree(tree).with_segments(segmentation)
} else if method == "random_forest" {
let classifier = RandomForest::new(X, control);
let gain = ClassifierGain { classifier };
let optimizer = TwoStepSearch { gain };
let mut segmentation = Segmentation::new(segmentation_type_enum, &optimizer);
tree = BinarySegmentationTree::new(X);
tree.grow(&mut segmentation);
BinarySegmentationResult::from_tree(tree).with_segments(segmentation)
} else if method == "change_in_mean" {
let gain = ChangeInMean::new(X, control);
let optimizer = GridSearch { gain };
let mut segmentation = Segmentation::new(segmentation_type_enum, &optimizer);
tree = BinarySegmentationTree::new(X);
tree.grow(&mut segmentation);
BinarySegmentationResult::from_tree(tree).with_segments(segmentation)
} else {
panic!("method should be one of 'knn', 'random_forest' or 'change_in_mean'. Got {method}",);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing;
use rstest::*;
#[rstest]
#[case("knn", "bs")]
#[case("knn", "wbs")]
#[case("knn", "sbs")]
#[case("change_in_mean", "bs")]
#[case("change_in_mean", "wbs")]
#[case("change_in_mean", "sbs")]
#[case("random_forest", "bs")]
#[case("random_forest", "sbs")]
fn test_binary_segmentation_wrapper(#[case] method: &str, #[case] segmentation_type: &str) {
let X = testing::array();
let control = Control::default().with_minimal_relative_segment_length(0.1);
assert_eq!(X.shape(), &[100, 5]);
assert_eq!(
changeforest(&X.view(), method, segmentation_type, &control).split_points(),
vec![25, 40, 80]
);
}
}