1use crate::algorithms::depth2::ConTreeDepth2;
2use crate::algorithms::interval_pruner::{Bound, IntervalsPruner};
3use crate::algorithms::shared;
4use crate::algorithms::support_feasible_splits;
5use crate::caching::{Cache, Entry};
6use crate::common::{
7 classification_error, FitOutcome, PointSelector, SearchConfig, SearchError, SearchStatus,
8 Statistics,
9};
10use crate::data::view::DataView;
11use crate::data::Dataset;
12use crate::tree::Tree;
13use rand::rngs::StdRng;
14use rand::SeedableRng;
15use std::collections::VecDeque;
16use std::time::Instant;
17
18pub struct ConTree {
24 config: SearchConfig,
25 statistics: Statistics,
26 cache: Cache,
27 specialized: ConTreeDepth2,
28 runtime: Instant,
29 rng: StdRng,
30 status: SearchStatus,
31 trajectory: Vec<(f64, usize)>,
33 pub tree: Tree,
34}
35
36impl ConTree {
37 pub fn with_config(config: SearchConfig) -> Self {
42 let mut solver = Self::new(
43 config.min_sup,
44 config.max_depth,
45 config.max_time,
46 config.max_error,
47 config.point_selector,
48 config.max_gap,
49 config.use_heuristic,
50 config.fast_d2,
51 );
52 solver.config = config;
53 solver
54 }
55
56 #[allow(clippy::too_many_arguments)]
58 pub fn new(
59 min_sup: usize,
60 max_depth: usize,
61 max_time: f64,
62 max_error: usize,
63 split_selection_strategy: PointSelector,
64 max_gap: usize,
65 use_heuristic: bool,
66 fast_d2: bool,
67 ) -> Self {
68 Self {
69 cache: Cache::default(),
70 config: SearchConfig::new(
71 min_sup,
72 max_depth,
73 max_time,
74 max_gap,
75 max_error,
76 use_heuristic,
77 fast_d2,
78 split_selection_strategy,
79 ),
80 statistics: Statistics::default(),
81 specialized: ConTreeDepth2,
82 runtime: Instant::now(),
83 rng: StdRng::from_os_rng(),
84 status: SearchStatus::Optimal,
85 trajectory: Vec::new(),
86 tree: Tree::default(),
87 }
88 }
89
90 pub fn with_random_state(mut self, seed: u64) -> Self {
93 self.rng = StdRng::seed_from_u64(seed);
94 self
95 }
96
97 pub fn fit(&mut self, dataset: &Dataset) -> Result<FitOutcome, SearchError> {
103 crate::algorithms::validate(&self.config, dataset)?;
104
105 let root_view = DataView::root(dataset, self.config.use_heuristic);
106
107 self.cache = Cache::new(self.config.max_depth, root_view.total_instances);
108
109 let root_index = self.cache.init();
110
111 self.statistics.num_features = root_view.get_feature_number();
112 self.statistics.num_samples = root_view.total_instances;
113
114 let (error, label) = classification_error(root_view.get_labels_freqs());
115 let mut entry = Entry {
116 error,
117 label,
118 ..Entry::default()
119 };
120
121 if let Some(entry) = self.cache.get_mut(root_index) {
122 self.config.max_error = self.config.max_error.min(error);
123 entry.error = error;
124 entry.label = label;
125 }
126
127 let root_config = self.config;
128 self.runtime = Instant::now();
129 self.trajectory.clear();
130 self.expand_node_with_view(
131 &root_view,
132 &root_config,
133 &mut entry,
134 0,
135 true,
136 root_config.max_error,
137 );
138
139 self.statistics.error = entry.error;
140 self.statistics.cache_size = self.cache.len();
141 self.statistics.duration = self.elapsed_time();
142 self.get_solution_tree();
143
144 self.status = if !self.time_remains() {
145 SearchStatus::TimeLimit
146 } else if entry.error <= self.config.max_gap && self.config.max_gap > 0 {
147 SearchStatus::ErrorBoundReached
148 } else {
149 SearchStatus::Optimal
150 };
151 self.tree.validate()?;
152
153 Ok(FitOutcome {
154 tree: self.tree.clone(),
155 statistics: self.statistics,
156 status: self.status,
157 })
158 }
159
160 pub fn trajectory(&self) -> &[(f64, usize)] {
163 &self.trajectory
164 }
165
166 pub fn status(&self) -> SearchStatus {
168 self.status
169 }
170
171 fn cache_child(&mut self, view: &DataView<'_>, child_depth: usize) -> (bool, usize, Entry) {
177 let (is_new, index) = self.cache.insert(&view.bitset, child_depth);
178 let depth = self.config.max_depth - child_depth;
179 let mut entry = Entry::default();
180
181 if let Some(slot) = self.cache.get_mut(index) {
182 if is_new {
183 let (error, label) = classification_error(view.get_labels_freqs());
184 slot.error = error;
185 slot.label = label;
186 slot.depth = depth;
187 slot.lower_bound = 0;
188 slot.is_valid = true;
189 }
190 entry = *slot;
191 }
192 (is_new, index, entry)
193 }
194
195 fn expand_node_with_view(
196 &mut self,
197 view: &DataView<'_>,
198 config: &SearchConfig,
199 current_best: &mut Entry,
200 parent_index: usize,
201 is_new: bool,
202 upper_bound: usize,
203 ) {
204 #[cfg(feature = "profiling")]
205 coz::progress!();
206 if current_best.error == 0 || view.is_empty() {
207 if let Some(entry) = self.cache.get_mut(parent_index) {
208 entry.is_leaf = true;
209 entry.mark_exact();
210 *current_best = *entry;
211 }
212 return;
213 }
214
215 if !is_new && current_best.is_optimal {
219 self.statistics.cache_hits += 1;
220 return;
221 }
222
223 if config.max_depth == 0 {
224 current_best.is_leaf = true;
225 current_best.mark_exact();
226
227 if let Some(entry) = self.cache.get_mut(parent_index) {
228 entry.is_leaf = true;
229 entry.mark_exact();
230 *current_best = *entry;
231 }
232
233 return;
234 }
235
236 if current_best.error <= config.max_gap || view.len() <= 1 {
237 current_best.mark_exact();
240 if let Some(entry) = self.cache.get_mut(parent_index) {
241 entry.mark_exact();
242 }
243 return;
244 }
245
246 if config.fast_d2 && config.max_depth == 2 {
249 let tree = self.specialized.fit(
250 view,
251 config,
252 current_best,
253 upper_bound,
254 &mut self.statistics,
255 );
256 current_best.finalize_lower_bound(upper_bound);
257 current_best.is_optimal = current_best.is_valid;
258 let tree_index = self.cache.insert_tree(tree);
259 if let Some(entry) = self.cache.get_mut(parent_index) {
260 *entry = *current_best;
261 entry.tree_idx = Some(tree_index);
262 }
263 self.statistics.specialized_solver_call += 1;
264
265 return;
266 }
267
268 let num_features = view.get_feature_number();
269 let heuristics_data = view.features_best_score();
270 debug_assert!(
271 num_features == heuristics_data.len(),
272 "Missmatch with heuristics and features number"
273 );
274 for &(_, feat) in heuristics_data.iter().take(num_features) {
275 self.expand_on_feature(
276 view,
277 feat,
278 parent_index,
279 config,
280 current_best,
281 upper_bound.min(current_best.error),
282 );
283
284 if current_best.error == 0 {
285 current_best.mark_exact();
286 if let Some(entry) = self.cache.get_mut(parent_index) {
287 entry.mark_exact();
288 }
289
290 return;
291 }
292 if !self.time_remains() {
293 return;
294 }
295 }
296
297 current_best.finalize_lower_bound(upper_bound);
301 current_best.is_optimal = current_best.is_valid;
302 if let Some(entry) = self.cache.get_mut(parent_index) {
303 entry.lower_bound = current_best.lower_bound;
304 entry.is_valid = current_best.is_valid;
305 entry.is_optimal = current_best.is_valid;
306 }
307 }
308
309 fn expand_on_feature(
310 &mut self,
311 view: &DataView<'_>,
312 feature_index: usize,
313 cache_index: usize,
314 config: &SearchConfig,
315 current_best: &mut Entry,
316 upper_bound: usize,
317 ) {
318 #[cfg(feature = "profiling")]
319 coz::scope!("expand_on_feature");
320 let feature_column = view.get_sorted_feature(feature_index);
321 let feature_column_ids = view.get_feature_indices(feature_index);
322
323 if config.point_selector == PointSelector::First {
324 self.expand_on_feature_gini_priority(
325 view,
326 feature_index,
327 cache_index,
328 config,
329 current_best,
330 upper_bound,
331 );
332 return;
333 }
334
335 let possible_index_split = view.get_possible_split_indices(feature_index);
336 if possible_index_split.is_empty() {
337 return;
338 }
339
340 let feasible =
344 support_feasible_splits(possible_index_split, view.len(), self.config.min_sup);
345 if feasible.is_empty() {
346 return;
347 }
348
349 let mut pruner = IntervalsPruner::new(possible_index_split, config.max_gap, config.min_sup);
350 let mut queue = VecDeque::new();
351 let init_bound = Bound::new(feasible.start, feasible.end - 1, None, None);
352 queue.push_back(init_bound);
353
354 while let Some(mut current_bound) = queue.pop_front() {
355 if !self.time_remains() {
356 return;
357 }
358
359 if pruner.subinterval_pruning(¤t_bound, current_best.error.min(upper_bound)) {
363 continue;
364 }
365
366 pruner.interval_shrinking(&mut current_bound, current_best.error.min(upper_bound));
367 if !current_bound.is_valid() {
368 continue;
369 }
370
371 let selected_point = self.select_point(config, ¤t_bound);
372 let split_point = possible_index_split[selected_point];
373 let int_half_distance = split_point
374 .saturating_sub(possible_index_split[current_bound.left_bound])
375 .max(possible_index_split[current_bound.right_bound].saturating_sub(split_point));
376
377 let threshold_value = if selected_point > 0 {
378 let previous = feature_column_ids[possible_index_split[selected_point - 1]];
379 let point = feature_column_ids[split_point];
380 shared::threshold_between(
381 feature_column[previous].value(),
382 feature_column[point].value(),
383 )
384 } else {
385 let point = feature_column_ids[split_point];
386 shared::threshold_between(
387 feature_column[feature_column_ids[0]].value(),
388 feature_column[point].value(),
389 )
390 };
391
392 let (left_view, right_view) = view.split(feature_index, split_point);
393
394 if left_view.len() < self.config.min_sup || right_view.len() < self.config.min_sup {
395 continue;
396 }
397
398 let process_left_first = left_view.len() >= right_view.len();
401
402 let left_config = config.derive_left();
403 let mut left_entry = Entry::default();
404 let mut right_entry = Entry::default();
405
406 let (mut left_index, mut right_index) = (0, 0);
409 let (mut left_is_new, mut right_is_new);
410
411 let larger_upper_bound = current_best.error.min(upper_bound.saturating_add(1));
412 self.statistics.general_solver_call += 1;
413
414 if process_left_first {
415 (left_is_new, left_index, left_entry) =
416 self.cache_child(&left_view, left_config.max_depth);
417
418 self.expand_node_with_view(
419 &left_view,
420 &left_config,
421 &mut left_entry,
422 left_index,
423 left_is_new,
424 larger_upper_bound,
425 );
426 left_entry.finalize_lower_bound(larger_upper_bound);
427 } else {
428 (right_is_new, right_index, right_entry) =
429 self.cache_child(&right_view, left_config.max_depth);
430
431 self.expand_node_with_view(
432 &right_view,
433 &left_config,
434 &mut right_entry,
435 right_index,
436 right_is_new,
437 larger_upper_bound,
438 );
439 right_entry.finalize_lower_bound(larger_upper_bound);
440 }
441
442 let larger_error = if process_left_first {
445 left_entry.lower_bound
446 } else {
447 right_entry.lower_bound
448 };
449 let budget =
452 current_best.error.min(upper_bound.saturating_add(1)) as i64 - larger_error as i64;
453 let smaller_ub = budget + int_half_distance as i64;
454 let smaller_upper_bound = smaller_ub.max(0) as usize;
455 let mut right_error: Option<usize> = None;
458
459 if smaller_ub > 0 || budget == 0 {
463 self.statistics.general_solver_call += 1;
464 let right_config = config.derive_right(left_config.max_gap);
465
466 if process_left_first {
467 (right_is_new, right_index, right_entry) =
468 self.cache_child(&right_view, right_config.max_depth);
469
470 self.expand_node_with_view(
471 &right_view,
472 &right_config,
473 &mut right_entry,
474 right_index,
475 right_is_new,
476 smaller_upper_bound,
477 );
478 right_entry.finalize_lower_bound(smaller_upper_bound);
479 } else {
480 (left_is_new, left_index, left_entry) =
481 self.cache_child(&left_view, right_config.max_depth);
482
483 self.expand_node_with_view(
484 &left_view,
485 &right_config,
486 &mut left_entry,
487 left_index,
488 left_is_new,
489 smaller_upper_bound,
490 );
491 left_entry.finalize_lower_bound(smaller_upper_bound);
492 }
493
494 right_error = Some(right_entry.lower_bound);
497
498 let feature_best = left_entry.lower_bound + right_entry.lower_bound;
499 if left_entry.is_valid && right_entry.is_valid && feature_best < current_best.error
500 {
501 current_best.error = feature_best;
502 if config.is_root {
503 self.trajectory.push((self.elapsed_time(), feature_best));
504 }
505 current_best.feature = feature_index;
506 current_best.split = threshold_value;
507 current_best.left = left_index;
508 current_best.right = right_index;
509
510 let is_optimal = feature_best == 0;
511 current_best.is_optimal = is_optimal;
512
513 if let Some(entry) = self.cache.get_mut(cache_index) {
514 *entry = *current_best;
515 }
516
517 if feature_best == 0 {
518 return;
519 }
520 }
521 }
522
523 let left_score = left_entry.is_valid.then_some(left_entry.lower_bound);
527 pruner.add_result(selected_point, left_score, right_error);
528 if current_bound.left_bound == current_bound.right_bound {
529 continue;
530 }
531
532 let score_difference = left_entry
533 .lower_bound
534 .saturating_add(right_error.unwrap_or(0))
535 .saturating_sub(current_best.error.min(upper_bound));
536
537 let (left_bound, right_bound) = pruner.neighbourhood_pruning(
538 score_difference,
539 current_bound.left_bound,
540 current_bound.right_bound,
541 selected_point,
542 );
543
544 if left_bound <= current_bound.right_bound {
545 queue.push_back(Bound {
546 left_bound,
547 right_bound: current_bound.right_bound,
548 last_split_left_index: Some(selected_point),
549 last_split_right_index: current_bound.last_split_right_index,
550 });
551 }
552
553 if current_bound.left_bound <= right_bound {
554 queue.push_back(Bound {
555 left_bound: current_bound.left_bound,
556 right_bound,
557 last_split_left_index: current_bound.last_split_left_index,
558 last_split_right_index: Some(selected_point),
559 });
560 }
561 }
562
563 if let Some(entry) = self.cache.get_mut(cache_index) {
564 entry.is_optimal = true;
565 }
566 }
567
568 fn expand_on_feature_gini_priority(
569 &mut self,
570 view: &DataView<'_>,
571 feature_index: usize,
572 cache_index: usize,
573 config: &SearchConfig,
574 current_best: &mut Entry,
575 upper_bound: usize,
576 ) {
577 let feature_column = view.get_sorted_feature(feature_index);
578 let feature_column_ids = view.get_feature_indices(feature_index);
579
580 let possible_splits = view.get_possible_split_indices(feature_index);
581
582 if possible_splits.is_empty() {
583 return;
584 }
585
586 let sorted_by_heuristic_indices = view.ordered_possible_splits(feature_index);
587
588 let mut pruner = IntervalsPruner::new(possible_splits, config.max_gap, config.min_sup);
589
590 let mut queue = VecDeque::new();
591 let init_bound = Bound::new(0, possible_splits.len() - 1, None, None);
592 queue.push_back(init_bound);
593
594 let mut pruned = vec![false; possible_splits.len()];
595
596 for &split_idx in sorted_by_heuristic_indices {
597 if !self.time_remains() {
598 return;
599 }
600
601 if pruned[split_idx] {
602 continue;
603 }
604
605 let mut current_left = split_idx;
608 while current_left > 0 && pruned[current_left - 1] {
609 current_left -= 1;
610 }
611 current_left = current_left.saturating_sub(1);
612
613 let mut current_right = split_idx;
614 while current_right < possible_splits.len() - 1 && pruned[current_right + 1] {
615 current_right += 1;
616 }
617 if current_right < possible_splits.len() - 1 {
618 current_right += 1;
619 }
620
621 let split_point = possible_splits[split_idx];
622
623 let threshold_value = if split_idx > 0 {
624 let previous = feature_column_ids[possible_splits[split_idx - 1]];
625 let point = feature_column_ids[split_point];
626 shared::threshold_between(
627 feature_column[previous].value(),
628 feature_column[point].value(),
629 )
630 } else {
631 let point = feature_column_ids[split_point];
632 shared::threshold_between(
633 feature_column[feature_column_ids[0]].value(),
634 feature_column[point].value(),
635 )
636 };
637
638 let (left_view, right_view) = view.split(feature_index, split_point);
639
640 if left_view.len() < self.config.min_sup || right_view.len() < self.config.min_sup {
641 pruned[split_idx] = true;
642 continue;
643 }
644
645 let process_left_first = left_view.len() >= right_view.len();
648
649 let left_config = config.derive_left();
650 let mut left_entry = Entry::default();
651 let mut right_entry = Entry::default();
652
653 let (mut left_index, mut right_index) = (0, 0);
656 let (mut left_is_new, mut right_is_new);
657
658 let int_half_distance = split_point
659 .saturating_sub(possible_splits[0])
660 .max(possible_splits[possible_splits.len() - 1].saturating_sub(split_point));
661
662 let larger_upper_bound = current_best.error.min(upper_bound.saturating_add(1));
663 self.statistics.general_solver_call += 1;
664
665 if process_left_first {
666 (left_is_new, left_index, left_entry) =
667 self.cache_child(&left_view, left_config.max_depth);
668
669 self.expand_node_with_view(
670 &left_view,
671 &left_config,
672 &mut left_entry,
673 left_index,
674 left_is_new,
675 larger_upper_bound,
676 );
677 left_entry.finalize_lower_bound(larger_upper_bound);
678 } else {
679 (right_is_new, right_index, right_entry) =
680 self.cache_child(&right_view, left_config.max_depth);
681
682 self.expand_node_with_view(
683 &right_view,
684 &left_config,
685 &mut right_entry,
686 right_index,
687 right_is_new,
688 larger_upper_bound,
689 );
690 right_entry.finalize_lower_bound(larger_upper_bound);
691 }
692
693 let larger_error = if process_left_first {
696 left_entry.lower_bound
697 } else {
698 right_entry.lower_bound
699 };
700 let budget =
703 current_best.error.min(upper_bound.saturating_add(1)) as i64 - larger_error as i64;
704 let smaller_ub = budget + int_half_distance as i64;
705 let smaller_upper_bound = smaller_ub.max(0) as usize;
706 let mut right_error: Option<usize> = None;
707
708 if smaller_ub > 0 || budget == 0 {
712 self.statistics.general_solver_call += 1;
713 let right_config = config.derive_right(left_config.max_gap);
714
715 if process_left_first {
716 (right_is_new, right_index, right_entry) =
717 self.cache_child(&right_view, right_config.max_depth);
718
719 self.expand_node_with_view(
720 &right_view,
721 &right_config,
722 &mut right_entry,
723 right_index,
724 right_is_new,
725 smaller_upper_bound,
726 );
727 right_entry.finalize_lower_bound(smaller_upper_bound);
728 } else {
729 (left_is_new, left_index, left_entry) =
730 self.cache_child(&left_view, right_config.max_depth);
731
732 self.expand_node_with_view(
733 &left_view,
734 &right_config,
735 &mut left_entry,
736 left_index,
737 left_is_new,
738 smaller_upper_bound,
739 );
740 left_entry.finalize_lower_bound(smaller_upper_bound);
741 }
742
743 right_error = Some(right_entry.lower_bound);
746 let feature_best = left_entry.lower_bound + right_entry.lower_bound;
747 if left_entry.is_valid && right_entry.is_valid && feature_best < current_best.error
748 {
749 current_best.error = feature_best;
750 if config.is_root {
751 self.trajectory.push((self.elapsed_time(), feature_best));
752 }
753 current_best.feature = feature_index;
754 current_best.split = threshold_value;
755 current_best.left = left_index;
756 current_best.right = right_index;
757
758 let is_optimal = feature_best == 0;
759 current_best.is_optimal = is_optimal;
760
761 if let Some(entry) = self.cache.get_mut(cache_index) {
762 *entry = *current_best;
763 }
764 }
765 }
766
767 let left_score = left_entry.is_valid.then_some(left_entry.lower_bound);
771 pruner.add_result(split_idx, left_score, right_error);
772
773 let score_difference = left_entry
774 .lower_bound
775 .saturating_add(right_error.unwrap_or(0))
776 .saturating_sub(current_best.error.min(upper_bound));
777 let (new_left_bound, new_right_bound) = pruner.neighbourhood_pruning(
778 score_difference,
779 0,
780 possible_splits.len() - 1,
781 split_idx,
782 );
783
784 let skip_from = new_right_bound.saturating_add(1).max(current_left);
788 if skip_from < split_idx {
789 pruned[skip_from..split_idx].fill(true);
790 }
791 let skip_to = new_left_bound.min(current_right + 1);
792 if split_idx + 1 < skip_to {
793 pruned[split_idx + 1..skip_to].fill(true);
794 }
795
796 if current_best.error == 0 {
797 break;
798 }
799 }
800
801 if let Some(entry) = self.cache.get_mut(cache_index) {
802 entry.is_optimal = true;
803 }
804 }
805
806 pub fn statistics(&self) -> Statistics {
807 self.statistics
808 }
809
810 fn time_remains(&self) -> bool {
811 shared::time_remains(&self.runtime, self.config.max_time)
812 }
813
814 fn elapsed_time(&self) -> f64 {
815 shared::elapsed_time(&self.runtime)
816 }
817
818 pub fn select_point(&mut self, config: &SearchConfig, bound: &Bound) -> usize {
819 shared::select_point(config.point_selector, &mut self.rng, bound)
820 }
821
822 pub fn get_solution_tree(&mut self) {
824 self.tree = shared::build_solution_tree(&self.cache);
825 }
826}
827
828#[cfg(test)]
829mod contree_test {
830 use crate::algorithms::continuous_tree::ConTree;
831 use crate::common::PointSelector;
832 use crate::reader::data_reader::DataReader;
833 use crate::reader::DataReaderError;
834 use crate::tests::fixture;
835
836 #[test]
837 fn deeper_search_never_scores_worse() -> Result<(), DataReaderError> {
838 let reader = DataReader::default();
839 let mut dataset = reader.read_file(&fixture("avila_1k.txt"))?;
840 dataset.sort_features();
841
842 let mut previous = usize::MAX;
843 for depth in 1..=3 {
844 let mut contree: ConTree = ConTree::new(
845 1,
846 depth,
847 100.0,
848 usize::MAX,
849 PointSelector::Mid,
850 0,
851 false,
852 true,
853 );
854 contree.fit(&dataset).expect("fit");
855
856 let error = contree.statistics().error;
857 assert!(
858 error <= previous,
859 "depth {depth} scored {error}, worse than depth {} at {previous}",
860 depth - 1
861 );
862 assert!(
863 error < dataset.count(),
864 "the search reported a nonsense error"
865 );
866 previous = error;
867 }
868
869 Ok(())
870 }
871
872 #[test]
873 fn the_depth_two_specialization_agrees_with_the_general_search() -> Result<(), DataReaderError>
874 {
875 let reader = DataReader::default();
876 let mut dataset = reader.read_file(&fixture("hepatitis.txt"))?;
877 dataset.sort_features();
878
879 let mut general: ConTree =
882 ConTree::new(1, 2, 100.0, usize::MAX, PointSelector::Mid, 0, false, false);
883 general.fit(&dataset).expect("fit");
884
885 let mut specialized: ConTree =
886 ConTree::new(1, 2, 100.0, usize::MAX, PointSelector::Mid, 0, false, true);
887 specialized.fit(&dataset).expect("fit");
888
889 assert_eq!(general.statistics().error, specialized.statistics().error);
890
891 Ok(())
892 }
893}