Skip to main content

contree/data/
view.rs

1use crate::bitsets::{BitCollection, Bitset, BitsetInit};
2use crate::data::{Dataset, Feature};
3
4/// Per-feature Gini scores used to order features and split candidates when
5/// the heuristic ordering is enabled.
6#[derive(Clone, Debug)]
7pub struct HeuristicValues {
8    /// For each feature, indices into its possible splits, best Gini first.
9    gini_per_split: Vec<Vec<usize>>,
10    /// `(best Gini, feature index)` for each feature; sorted best first by
11    /// [`Self::sort_by_gini`].
12    best_gini_per_feature: Vec<(f64, usize)>,
13}
14
15impl HeuristicValues {
16    /// Neutral scores (Gini 1.0) for `num_features` features, in index order.
17    pub fn new(num_features: usize) -> Self {
18        Self {
19            gini_per_split: vec![Vec::new(); num_features],
20            best_gini_per_feature: (0..num_features).map(|i| (1.0, i)).collect(),
21        }
22    }
23
24    /// Stores the split order and best Gini of `feature`. `split_indices` must
25    /// already be sorted best first.
26    pub fn set_feature_ginis(&mut self, feature: usize, split_indices: Vec<usize>, best_gini: f64) {
27        if feature < self.gini_per_split.len() {
28            self.gini_per_split[feature] = split_indices;
29            self.best_gini_per_feature[feature] = (best_gini, feature);
30        }
31    }
32
33    /// Sorts features by their best Gini, lowest first.
34    pub fn sort_by_gini(&mut self) {
35        self.best_gini_per_feature
36            .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
37    }
38}
39
40/// A subset of the dataset: the instances that reach one node of the tree.
41///
42/// Each feature column is kept sorted by value, so candidate thresholds are
43/// the positions where the value changes. Splitting a view partitions every
44/// column in one linear pass, preserving the order.
45pub struct DataView<'a> {
46    /// The full dataset the view refers into.
47    pub dataset: &'a Dataset,
48    /// Number of instances in the full dataset.
49    pub total_instances: usize,
50    /// For each feature, positions into `dataset[f]` of the instances in the
51    /// view, in increasing value order.
52    pub feature_columns: Vec<Vec<usize>>,
53    /// For each feature, positions in `feature_columns[f]` where the value
54    /// changes: the candidate thresholds.
55    pub possible_split_indices: Vec<Vec<usize>>,
56    /// Number of instances of each class in the view.
57    pub label_freq: Vec<usize>,
58    /// Whether features and splits are ordered by Gini.
59    pub sort_by_heuristic: bool,
60    /// Gini ordering of features and splits, when `sort_by_heuristic` is set.
61    pub heuristic_values: HeuristicValues,
62    /// The instances in the view; the cache key of the subproblem.
63    pub bitset: Bitset,
64}
65
66impl<'a> DataView<'a> {
67    /// The view of the whole dataset. The dataset's features must be sorted.
68    pub fn root(dataset: &'a Dataset, sort_by_heuristic: bool) -> Self {
69        let total_instances = dataset.count();
70        let mut label_freq = vec![0; dataset.num_labels()];
71        let mut feature_columns = Vec::new();
72        let mut possible_split_indices = Vec::new();
73
74        for (feature_idx, feature) in dataset.into_iter().enumerate() {
75            let feature_len = feature.len();
76
77            let mut splits = Vec::new();
78            let mut last_unique_index: Option<usize> = None;
79
80            for pos in 0..feature_len {
81                let el = &feature[pos];
82
83                if feature_idx == 0 {
84                    label_freq[el.label as usize] += 1;
85                }
86
87                if let Some(last) = last_unique_index {
88                    if el.unique_value_idx != last {
89                        splits.push(pos);
90                    }
91                }
92                last_unique_index = Some(el.unique_value_idx);
93            }
94            possible_split_indices.push(splits);
95            feature_columns.push((0..feature_len).collect::<Vec<usize>>());
96        }
97
98        let mut heuristic_values = HeuristicValues::new(dataset.num_features());
99
100        if sort_by_heuristic {
101            for feature_idx in 0..dataset.num_features() {
102                let feature = &dataset[feature_idx];
103                let idx = &feature_columns[feature_idx];
104                let (ordered_index, gini) = Self::compute_gini_for_all_splits(
105                    feature,
106                    idx,
107                    &possible_split_indices[feature_idx],
108                    &label_freq,
109                    dataset.num_labels(),
110                );
111
112                heuristic_values.set_feature_ginis(feature_idx, ordered_index, gini);
113            }
114            heuristic_values.sort_by_gini();
115        }
116
117        let mut bitset = Bitset::new(BitsetInit::Full(total_instances));
118        bitset.save_count();
119
120        Self {
121            dataset,
122            total_instances,
123            feature_columns,
124            possible_split_indices,
125            label_freq,
126            sort_by_heuristic,
127            heuristic_values,
128            bitset,
129        }
130    }
131
132    /// Number of instances in the view.
133    pub fn get_dataset_size(&self) -> usize {
134        self.feature_columns[0].len()
135    }
136
137    /// The ids of the instances in the view.
138    pub fn instance_ids(&self) -> Vec<usize> {
139        let column = &self.dataset[0];
140        self.feature_columns[0]
141            .iter()
142            .map(|&pos| column[pos].tid)
143            .collect()
144    }
145
146    /// Number of features.
147    pub fn get_feature_number(&self) -> usize {
148        self.dataset.num_features()
149    }
150
151    /// The full, sorted column of feature `f`.
152    pub fn get_sorted_feature(&self, f: usize) -> &Feature {
153        &self.dataset[f]
154    }
155
156    /// Positions into [`Self::get_sorted_feature`] of the instances in the
157    /// view, in value order.
158    pub fn get_feature_indices(&self, f: usize) -> &[usize] {
159        &self.feature_columns[f]
160    }
161
162    /// Number of instances of each class in the view.
163    pub fn get_labels_freqs(&self) -> &[usize] {
164        &self.label_freq
165    }
166
167    /// Number of classes in the dataset.
168    pub fn get_num_labels(&self) -> usize {
169        self.dataset.num_labels()
170    }
171
172    /// Candidate thresholds of feature `f`, as positions in the sorted column.
173    pub fn get_possible_split_indices(&self, f: usize) -> &[usize] {
174        &self.possible_split_indices[f]
175    }
176
177    /// Largest number of candidate thresholds of any feature.
178    pub fn get_max_splits(&self) -> usize {
179        self.possible_split_indices
180            .iter()
181            .map(|x| x.len())
182            .max()
183            .unwrap_or(0)
184    }
185
186    /// Indices into the candidate thresholds of `feature`, best Gini first.
187    /// Empty unless the heuristic ordering is enabled.
188    pub fn ordered_possible_splits(&self, feature: usize) -> &[usize] {
189        &self.heuristic_values.gini_per_split[feature]
190    }
191
192    /// `(best Gini, feature)` pairs, best first when the heuristic ordering is
193    /// enabled and in feature order otherwise.
194    pub fn features_best_score(&self) -> &[(f64, usize)] {
195        &self.heuristic_values.best_gini_per_feature
196    }
197
198    #[inline]
199    fn compute_split_indices_for(col: &Feature, idxs: &[usize]) -> Vec<usize> {
200        if idxs.is_empty() {
201            return Vec::new();
202        }
203        let mut out = Vec::with_capacity(idxs.len() / 10);
204        let mut last = None::<usize>;
205
206        for (i, &pos) in idxs.iter().enumerate() {
207            let cur = col[pos].unique_value_idx;
208            if i > 0 && Some(cur) != last {
209                out.push(i);
210            }
211            last = Some(cur);
212        }
213
214        out.shrink_to_fit();
215        out
216    }
217
218    /// Fills the class histograms of both sides of a split of `feature_index`
219    /// at `split_point`. Both slices must start zeroed.
220    pub fn initialize_split_parameters(
221        &self,
222        feature_index: usize,
223        split_point: usize,
224        left_freq: &mut [usize],
225        right_freq: &mut [usize],
226    ) {
227        let num_labels = self.dataset.num_labels();
228
229        let total_size = self.get_dataset_size();
230        let feature_ids = self.get_feature_indices(feature_index);
231        let feature = self.get_sorted_feature(feature_index);
232
233        // Count the smaller side and derive the other by subtraction.
234        if 2 * split_point < total_size {
235            for i in 0..split_point {
236                let data_point = &feature[feature_ids[i]];
237                left_freq[data_point.label as usize] += 1;
238            }
239            for label in 0..num_labels {
240                right_freq[label] = self.label_freq[label] - left_freq[label];
241            }
242        } else {
243            for i in split_point..total_size {
244                let data_point = &feature[feature_ids[i]];
245                right_freq[data_point.label as usize] += 1;
246            }
247            for label in 0..num_labels {
248                left_freq[label] = self.label_freq[label] - right_freq[label];
249            }
250        }
251    }
252
253    /// Computes the weighted Gini of every candidate threshold of a feature.
254    ///
255    /// Returns the indices of the candidates sorted best first, and the best
256    /// Gini value.
257    fn compute_gini_for_all_splits(
258        feature: &Feature,
259        idx: &[usize],
260        possible_splits: &[usize],
261        label_freq: &[usize],
262        num_labels: usize,
263    ) -> (Vec<usize>, f64) {
264        if feature.is_empty() || possible_splits.is_empty() {
265            return (Vec::new(), 1.0);
266        }
267
268        let mut gini_values = Vec::with_capacity(possible_splits.len());
269        let mut left_label_freq = vec![0usize; num_labels];
270        let mut right_label_freq = label_freq.to_vec();
271        let mut best_gini = 1.0;
272
273        let mut last_pos = 0;
274
275        for (split_idx, &split_pos) in possible_splits.iter().enumerate() {
276            // Move the instances before this threshold to the left side.
277            for &data_idx in &idx[last_pos..split_pos] {
278                let label = feature[data_idx].label as usize;
279                right_label_freq[label] -= 1;
280                left_label_freq[label] += 1;
281            }
282
283            let left_count = split_pos;
284            let right_count = idx.len() - left_count;
285
286            let gini = Self::compute_gini_from_frequencies(
287                &left_label_freq,
288                &right_label_freq,
289                left_count,
290                right_count,
291            );
292
293            if gini < best_gini {
294                best_gini = gini;
295            }
296
297            gini_values.push((split_idx, gini));
298            last_pos = split_pos;
299        }
300
301        gini_values.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
302
303        let sorted_indices: Vec<usize> = gini_values.into_iter().map(|(idx, _)| idx).collect();
304
305        (sorted_indices, best_gini)
306    }
307
308    /// Weighted Gini impurity of a split, from the class histograms of its
309    /// two sides.
310    #[inline]
311    fn compute_gini_from_frequencies(
312        left_freq: &[usize],
313        right_freq: &[usize],
314        left_count: usize,
315        right_count: usize,
316    ) -> f64 {
317        let mut left_gini = 1.0;
318        let mut right_gini = 1.0;
319
320        if left_count > 0 {
321            for &freq in left_freq {
322                let prob = freq as f64 / left_count as f64;
323                left_gini -= prob * prob;
324            }
325        }
326
327        if right_count > 0 {
328            for &freq in right_freq {
329                let prob = freq as f64 / right_count as f64;
330                right_gini -= prob * prob;
331            }
332        }
333
334        let total = left_count + right_count;
335        if total > 0 {
336            (left_gini * left_count as f64 + right_gini * right_count as f64) / total as f64
337        } else {
338            1.0
339        }
340    }
341
342    /// Splits the view on feature `sf` at position `split_point` of its sorted
343    /// column. The left view holds the instances before the position.
344    pub fn split(&self, sf: usize, split_point: usize) -> (Self, Self) {
345        #[cfg(feature = "profiling")]
346        coz::scope!("Split view");
347        let col = &self.dataset[sf];
348        let sf_idxs = &self.feature_columns[sf];
349
350        let (sf_left_idxs, sf_right_idxs) = sf_idxs.split_at(split_point);
351
352        let mut left_bitset = Bitset::new(BitsetInit::Empty(self.total_instances));
353        let mut left_label_freq = vec![0; self.dataset.num_labels()];
354
355        for &pos in sf_left_idxs {
356            let el = &col[pos];
357            left_bitset.set(el.tid);
358            left_label_freq[el.label as usize] += 1;
359        }
360        left_bitset.save_count();
361
362        let mut right_label_freq = vec![0; self.dataset.num_labels()];
363        for label in 0..self.dataset.num_labels() {
364            right_label_freq[label] = self.label_freq[label] - left_label_freq[label];
365        }
366
367        let num_features = self.get_feature_number();
368        let left_size_estimate = split_point;
369        let right_size_estimate = self.get_dataset_size() - split_point;
370
371        let mut left_pfi: Vec<Vec<usize>> = (0..num_features)
372            .map(|_| Vec::with_capacity(left_size_estimate))
373            .collect();
374        let mut right_pfi: Vec<Vec<usize>> = (0..num_features)
375            .map(|_| Vec::with_capacity(right_size_estimate))
376            .collect();
377
378        let mut left_split_indices: Vec<Vec<usize>> = (0..num_features)
379            .map(|_| Vec::with_capacity(left_size_estimate / 10))
380            .collect();
381        let mut right_split_indices: Vec<Vec<usize>> = (0..num_features)
382            .map(|_| Vec::with_capacity(right_size_estimate / 10))
383            .collect();
384
385        let mut left_heuristics = HeuristicValues::new(num_features);
386        let mut right_heuristics = HeuristicValues::new(num_features);
387
388        // Scratch buffers for the Gini pass, reused across features. They are
389        // empty when the heuristic ordering is off.
390        let num_labels = self.dataset.num_labels();
391        let heuristic_width = if self.sort_by_heuristic {
392            num_labels
393        } else {
394            0
395        };
396        let mut left_label_running = vec![0usize; heuristic_width];
397        let mut right_label_running = vec![0usize; heuristic_width];
398        let mut left_label_remaining = vec![0usize; heuristic_width];
399        let mut right_label_remaining = vec![0usize; heuristic_width];
400        let mut left_gini_values: Vec<(usize, f64)> = Vec::new();
401        let mut right_gini_values: Vec<(usize, f64)> = Vec::new();
402
403        for f in 0..num_features {
404            if f == sf {
405                // The split feature is already partitioned by `split_at`.
406                left_pfi[f] = sf_left_idxs.to_vec();
407                right_pfi[f] = sf_right_idxs.to_vec();
408
409                left_split_indices[f] =
410                    Self::compute_split_indices_for(&self.dataset[f], sf_left_idxs);
411                right_split_indices[f] =
412                    Self::compute_split_indices_for(&self.dataset[f], sf_right_idxs);
413
414                if self.sort_by_heuristic {
415                    let (ordered_splits, best_gini) = Self::compute_gini_for_all_splits(
416                        &self.dataset[f],
417                        &left_pfi[f],
418                        &left_split_indices[f],
419                        &left_label_freq,
420                        self.dataset.num_labels(),
421                    );
422                    left_heuristics.set_feature_ginis(f, ordered_splits, best_gini);
423
424                    let (ordered_splits, best_gini) = Self::compute_gini_for_all_splits(
425                        &self.dataset[f],
426                        &right_pfi[f],
427                        &right_split_indices[f],
428                        &right_label_freq,
429                        self.dataset.num_labels(),
430                    );
431                    right_heuristics.set_feature_ginis(f, ordered_splits, best_gini);
432                }
433                continue;
434            }
435
436            let fcol = &self.dataset[f];
437            let parent_idxs = &self.feature_columns[f];
438
439            let mut left_last_unique: Option<usize> = None;
440            let mut right_last_unique: Option<usize> = None;
441            let mut left_counter = 0;
442            let mut right_counter = 0;
443
444            let mut left_best_gini = 1.0;
445            let mut right_best_gini = 1.0;
446
447            if self.sort_by_heuristic {
448                left_gini_values.clear();
449                right_gini_values.clear();
450                left_label_running.fill(0);
451                right_label_running.fill(0);
452                left_label_remaining.copy_from_slice(&left_label_freq);
453                right_label_remaining.copy_from_slice(&right_label_freq);
454            }
455
456            for &pos in parent_idxs {
457                let el = &fcol[pos];
458                let row = el.tid;
459                let label = el.label as usize;
460
461                if left_bitset.contains(row) {
462                    left_pfi[f].push(pos);
463
464                    if let Some(last) = left_last_unique {
465                        if el.unique_value_idx != last {
466                            left_split_indices[f].push(left_counter);
467
468                            if self.sort_by_heuristic {
469                                let gini = Self::compute_gini_from_frequencies(
470                                    &left_label_running,
471                                    &left_label_remaining,
472                                    left_counter,
473                                    left_size_estimate - left_counter,
474                                );
475                                if gini < left_best_gini {
476                                    left_best_gini = gini;
477                                }
478                                left_gini_values.push((left_split_indices[f].len() - 1, gini));
479                            }
480                        }
481                    }
482                    left_last_unique = Some(el.unique_value_idx);
483                    left_counter += 1;
484
485                    if self.sort_by_heuristic {
486                        left_label_remaining[label] -= 1;
487                        left_label_running[label] += 1;
488                    }
489                } else {
490                    right_pfi[f].push(pos);
491
492                    if let Some(last) = right_last_unique {
493                        if el.unique_value_idx != last {
494                            right_split_indices[f].push(right_counter);
495
496                            if self.sort_by_heuristic {
497                                let gini = Self::compute_gini_from_frequencies(
498                                    &right_label_running,
499                                    &right_label_remaining,
500                                    right_counter,
501                                    right_size_estimate - right_counter,
502                                );
503                                if gini < right_best_gini {
504                                    right_best_gini = gini;
505                                }
506                                right_gini_values.push((right_split_indices[f].len() - 1, gini));
507                            }
508                        }
509                    }
510                    right_last_unique = Some(el.unique_value_idx);
511                    right_counter += 1;
512
513                    if self.sort_by_heuristic {
514                        right_label_remaining[label] -= 1;
515                        right_label_running[label] += 1;
516                    }
517                }
518            }
519
520            if self.sort_by_heuristic {
521                left_gini_values
522                    .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
523                right_gini_values
524                    .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
525
526                let left_sorted_indices: Vec<usize> =
527                    left_gini_values.iter().map(|&(idx, _)| idx).collect();
528                let right_sorted_indices: Vec<usize> =
529                    right_gini_values.iter().map(|&(idx, _)| idx).collect();
530
531                left_heuristics.set_feature_ginis(f, left_sorted_indices, left_best_gini);
532                right_heuristics.set_feature_ginis(f, right_sorted_indices, right_best_gini);
533            }
534
535            // No `shrink_to_fit`: a view only lives while its node is expanded.
536        }
537
538        let mut right_bitset = self.bitset.intersect_with(&left_bitset, true);
539        right_bitset.save_count();
540
541        if self.sort_by_heuristic {
542            left_heuristics.sort_by_gini();
543            right_heuristics.sort_by_gini();
544        }
545
546        let left = Self {
547            dataset: self.dataset,
548            total_instances: self.total_instances,
549            feature_columns: left_pfi,
550            possible_split_indices: left_split_indices,
551            label_freq: left_label_freq,
552            sort_by_heuristic: self.sort_by_heuristic,
553            heuristic_values: left_heuristics,
554            bitset: left_bitset,
555        };
556
557        let right = Self {
558            dataset: self.dataset,
559            total_instances: self.total_instances,
560            feature_columns: right_pfi,
561            possible_split_indices: right_split_indices,
562            label_freq: right_label_freq,
563            sort_by_heuristic: self.sort_by_heuristic,
564            heuristic_values: right_heuristics,
565            bitset: right_bitset,
566        };
567
568        (left, right)
569    }
570
571    /// Number of instances in the view.
572    pub fn len(&self) -> usize {
573        debug_assert!(
574            self.bitset.count() == self.feature_columns[0].len(),
575            "Mismatched size in bitset and feature columns"
576        );
577        self.feature_columns[0].len()
578    }
579
580    /// Whether the view has no instances.
581    pub fn is_empty(&self) -> bool {
582        self.len() == 0
583    }
584}
585
586#[cfg(test)]
587mod data_view_tests {
588    use crate::bitsets::BitCollection;
589    use crate::data::view::DataView;
590    use crate::reader::data_reader::DataReader;
591    use crate::reader::DataReaderError;
592    use crate::tests::fixture;
593
594    #[test]
595    fn root_view_covers_the_whole_dataset() -> Result<(), DataReaderError> {
596        let reader = DataReader::default();
597        let mut dataset = reader.read_file(&fixture("avila_1k.txt"))?;
598        dataset.sort_features();
599        dataset.compute_unique_feature_values();
600
601        let view = DataView::root(&dataset, true);
602
603        assert_eq!(view.len(), dataset.count());
604        assert_eq!(view.bitset.count(), dataset.count());
605        assert_eq!(view.get_feature_number(), dataset.num_features());
606        assert_eq!(
607            view.get_labels_freqs().iter().sum::<usize>(),
608            dataset.count(),
609            "the root label histogram must account for every instance"
610        );
611
612        // Split candidates are positions strictly inside the column, and are
613        // reported in increasing order.
614        for f in 0..dataset.num_features() {
615            let splits = view.get_possible_split_indices(f);
616            assert!(splits.windows(2).all(|w| w[0] < w[1]));
617            assert!(splits.iter().all(|&i| i > 0 && i < dataset.count()));
618        }
619
620        Ok(())
621    }
622
623    #[test]
624    fn split_partitions_the_view_without_loss() -> Result<(), DataReaderError> {
625        let reader = DataReader::default();
626        let mut dataset = reader.read_file(&fixture("avila_1k.txt"))?;
627        dataset.sort_features();
628        dataset.compute_unique_feature_values();
629
630        let view = DataView::root(&dataset, false);
631        let feature = 4;
632        let split = view.get_possible_split_indices(feature)[3];
633        let (left, right) = view.split(feature, split);
634
635        assert_eq!(left.len() + right.len(), view.len());
636        assert_eq!(left.bitset.count(), left.len());
637        assert_eq!(right.bitset.count(), right.len());
638        for label in 0..dataset.num_labels() {
639            assert_eq!(
640                left.get_labels_freqs()[label] + right.get_labels_freqs()[label],
641                view.get_labels_freqs()[label]
642            );
643        }
644
645        Ok(())
646    }
647}