Skip to main content

dtrees_rs/caching/
cover_map.rs

1use crate::caching::entry::{CacheEntry, CacheEntryUpdater};
2use crate::caching::helpers::{CacheKey, Index};
3use crate::caching::Caching;
4use crate::cover::Cover;
5use rustc_hash::FxHashMap;
6
7/// The subproblems of one depth and one number of instances, by instances.
8type Bucket = FxHashMap<Box<[u64]>, usize>;
9
10/// A [`Caching`] implementation keyed by the depth and the instances of a
11/// subproblem: paths reaching the same instances at the same depth share one
12/// entry.
13///
14/// Entries live in an arena, found by depth, then number of instances, then
15/// instances. Each entry keeps its instances, one word per 64 rows, so the
16/// memory grows with the size of the dataset.
17#[derive(Default)]
18pub struct CoverMap {
19    arena: Vec<CacheEntry>,
20    map: Vec<Vec<Bucket>>,
21    /// Words of the instances being inserted, reused between insertions.
22    buffer: Vec<u64>,
23    /// Number of instances of the dataset being fitted.
24    num_samples: Option<usize>,
25}
26
27impl Caching for CoverMap {
28    fn init(&mut self) -> Index {
29        self.arena.clear();
30        self.map.clear();
31        self.num_samples = None;
32        self.arena.push(CacheEntry::default());
33        Index::new_at(0)
34    }
35
36    fn root_index(&mut self) -> Index {
37        Index::new_at(0)
38    }
39
40    fn root(&self) -> Option<&CacheEntry> {
41        self.arena.first()
42    }
43
44    fn insert(&mut self, path: &[usize], cover: &Cover) -> Index {
45        let num_samples = *self.num_samples.get_or_insert(cover.num_samples);
46        debug_assert_eq!(num_samples, cover.num_samples, "one dataset per fit");
47        cover.write_words(&mut self.buffer);
48        let bucket = bucket(&mut self.map, path.len(), cover.count());
49        let index = match bucket.get(self.buffer.as_slice()) {
50            Some(&index) => index,
51            None => {
52                let index = self.arena.len();
53                bucket.insert(self.buffer.as_slice().into(), index);
54                self.arena.push(CacheEntry::default());
55                index
56            }
57        };
58        if self.arena[index].has_finite_leaf_error() {
59            Index::existing(index)
60        } else {
61            Index::new_at(index)
62        }
63    }
64
65    fn find(&self, path: &[usize], cover: &Cover) -> Option<&CacheEntry> {
66        if path.is_empty() {
67            debug_assert_eq!(
68                cover.count(),
69                cover.num_samples,
70                "the root has every instance"
71            );
72            return self.root();
73        }
74        let mut words = Vec::new();
75        cover.write_words(&mut words);
76        let index = self
77            .map
78            .get(path.len())?
79            .get(cover.count())?
80            .get(words.as_slice())?;
81        self.arena.get(*index)
82    }
83
84    fn keyed_by_cover(&self) -> bool {
85        true
86    }
87
88    fn node(&self, key: &CacheKey) -> Option<&CacheEntry> {
89        self.arena.get(key.0)
90    }
91
92    fn contains(&self, key: &CacheKey) -> bool {
93        self.node(key).is_some()
94    }
95
96    fn update_root(&mut self) -> Option<CacheEntryUpdater<'_>> {
97        self.arena.first_mut().map(CacheEntryUpdater::new)
98    }
99
100    fn update_node(&mut self, key: &CacheKey) -> Option<CacheEntryUpdater<'_>> {
101        self.arena.get_mut(key.0).map(CacheEntryUpdater::new)
102    }
103
104    fn size(&self) -> usize {
105        self.arena.len()
106    }
107
108    fn is_empty(&self) -> bool {
109        self.arena.is_empty()
110    }
111}
112
113/// The map of the subproblems at `depth` with `count` instances, created if
114/// needed.
115fn bucket(map: &mut Vec<Vec<Bucket>>, depth: usize, count: usize) -> &mut Bucket {
116    if map.len() <= depth {
117        map.resize_with(depth + 1, Vec::new);
118    }
119    let row = &mut map[depth];
120    if row.len() <= count {
121        row.resize_with(count + 1, FxHashMap::default);
122    }
123    &mut row[count]
124}
125
126#[cfg(test)]
127mod cover_map_test {
128    use crate::bitsets::{BitCollection, Bitset, BitsetInit};
129    use crate::caching::cover_map::CoverMap;
130    use crate::caching::Caching;
131    use crate::cover::Cover;
132    use crate::reader::data_reader::DataReader;
133    use rustc_hash::{FxHashMap, FxHashSet};
134    use std::hash::BuildHasher;
135    use std::path::Path;
136
137    #[test]
138    fn paths_reaching_the_same_instances_share_an_entry() {
139        // Features 0 and 1 are both 1 on instances 0 and 1, so items 1 and 3
140        // reach the same instances.
141        let mut feature = Bitset::new(BitsetInit::Empty(4));
142        feature.set(0);
143        feature.set(1);
144        let mut cover = Cover::new(
145            vec![feature.clone(), feature],
146            vec![Bitset::new(BitsetInit::Full(4))],
147            4,
148        );
149        let mut cache = CoverMap::default();
150        cache.init();
151
152        cover.branch_on(1);
153        let first = cache.insert(&[1], &cover);
154        cover.backtrack();
155        cover.branch_on(3);
156        let second = cache.insert(&[3], &cover);
157        assert!(first.is_new() && second.is_new());
158        assert_eq!(first.position(), second.position());
159        assert!(cache.find(&[1], &cover).is_some());
160
161        // The same instances one level deeper are another subproblem.
162        cover.branch_on(1);
163        let deeper = cache.insert(&[1, 3], &cover);
164        assert_ne!(deeper.position(), first.position());
165        assert_eq!(cache.size(), 3);
166    }
167
168    /// Adds the key of every subproblem up to `depth` below the current one,
169    /// testing features after `first_item`, to its `(depth, count)` bucket.
170    fn collect_keys(
171        cover: &mut Cover,
172        first_item: usize,
173        depth: usize,
174        keys: &mut FxHashMap<(usize, usize), FxHashSet<Box<[u64]>>>,
175    ) {
176        let mut words = Vec::new();
177        cover.write_words(&mut words);
178        keys.entry((depth, cover.count()))
179            .or_default()
180            .insert(words.into());
181        if depth == 0 {
182            return;
183        }
184        for item in first_item..2 * cover.num_attributes {
185            if cover.count_if_branch_on(item) > 0 {
186                cover.branch_on(item);
187                collect_keys(cover, (item / 2 + 1) * 2, depth - 1, keys);
188                cover.backtrack();
189            }
190        }
191    }
192
193    #[test]
194    fn distinct_subproblems_have_distinct_hashes() {
195        let mut cover = DataReader::default()
196            .read_file(Path::new("test_data/german-credit.txt"))
197            .expect("the dataset reads");
198        let mut keys = FxHashMap::default();
199        collect_keys(&mut cover, 0, 3, &mut keys);
200
201        let hasher = *FxHashMap::<Box<[u64]>, usize>::default().hasher();
202        let mut total = 0;
203        for bucket in keys.values() {
204            let hashes: FxHashSet<u64> = bucket.iter().map(|key| hasher.hash_one(key)).collect();
205            assert_eq!(hashes.len(), bucket.len());
206            total += bucket.len();
207        }
208        assert_eq!(total, 582_030);
209    }
210}