Skip to main content

dtrees_rs/algorithms/optimal/rules/
topk.rs

1use crate::algorithms::optimal::rules::core::Reason;
2use crate::algorithms::optimal::rules::helpers::StepStrategy;
3use crate::algorithms::optimal::rules::{Rule, RuleContext, RuleResult, RuleState};
4
5/// Top-k search: at each node, a pass only branches on the `k + 1`
6/// best-ranked features. Relaxing the rule raises `k`, following its
7/// [`StepStrategy`], up to `limit`.
8pub struct TopkRule {
9    limit: usize,
10    budget: usize,
11    increment: Box<dyn StepStrategy>,
12    priority: u8,
13    state: RuleState,
14    relaxable: bool,
15    delay: u8,
16}
17
18impl TopkRule {
19    /// A rule whose `k` grows following `increment`, up to `limit`.
20    pub fn new(limit: usize, increment: Box<dyn StepStrategy>) -> Self {
21        Self {
22            limit,
23            budget: 0,
24            increment,
25            priority: 90,
26            state: RuleState::Active,
27            relaxable: true,
28            delay: 0,
29        }
30    }
31
32    /// Sets the number of passes before the rule takes effect.
33    pub fn with_delay(mut self, delay: u8) -> Self {
34        self.delay = delay;
35        self
36    }
37
38    /// Sets `k` for the first pass.
39    pub fn with_budget(mut self, budget: usize) -> Self {
40        self.budget = budget;
41        self
42    }
43}
44
45impl Rule for TopkRule {
46    fn evaluate(&self, context: &RuleContext) -> RuleResult {
47        if context.position > self.budget {
48            RuleResult::stop_with_bound(f64::INFINITY, Reason::RuleReason)
49        } else {
50            RuleResult::continue_search()
51        }
52    }
53
54    fn priority(&self) -> u8 {
55        self.priority
56    }
57
58    fn description(&self) -> String {
59        "TopK rule".to_string()
60    }
61
62    fn state(&self) -> RuleState {
63        self.state
64    }
65
66    fn is_active(&self) -> bool {
67        self.state == RuleState::Active
68    }
69
70    fn activate(&mut self) {
71        self.state = RuleState::Active
72    }
73
74    fn deactivate(&mut self) {
75        self.state = RuleState::Disabled
76    }
77
78    fn relax(&mut self) {
79        if !self.is_active() {
80            return;
81        }
82
83        if self.is_relaxable() && self.budget >= self.limit {
84            self.deactivate();
85            return;
86        }
87        self.budget = self.increment.next();
88        if self.budget >= self.limit {
89            self.budget = self.limit;
90        }
91    }
92
93    fn is_relaxable(&self) -> bool {
94        self.relaxable
95    }
96
97    fn delay(&self) -> u8 {
98        self.delay
99    }
100
101    fn as_any(&self) -> &dyn std::any::Any {
102        self
103    }
104
105    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
106        self
107    }
108}
109
110/// Like [`TopkRule`], but the number of features allowed halves at each
111/// level: `k / 2^depth`, and at least one.
112pub struct DecreasingTopkRule {
113    limit: usize,
114    budget: usize,
115    increment: Box<dyn StepStrategy>,
116    priority: u8,
117    state: RuleState,
118    relaxable: bool,
119    delay: u8,
120}
121
122impl DecreasingTopkRule {
123    /// A rule whose root `k` grows following `increment`, up to `limit`.
124    pub fn new(limit: usize, increment: Box<dyn StepStrategy>) -> Self {
125        Self {
126            limit,
127            budget: 0,
128            increment,
129            priority: 90,
130            state: RuleState::Active,
131            relaxable: true,
132            delay: 0,
133        }
134    }
135
136    /// Sets the number of passes before the rule takes effect.
137    pub fn with_delay(mut self, delay: u8) -> Self {
138        self.delay = delay;
139        self
140    }
141
142    /// Sets `k` for the first pass.
143    pub fn with_budget(mut self, budget: usize) -> Self {
144        self.budget = budget;
145        self
146    }
147}
148
149impl Rule for DecreasingTopkRule {
150    fn evaluate(&self, context: &RuleContext) -> RuleResult {
151        let depth_budget = (self.budget / (2.0_f64.powi(context.depth as i32) as usize)).max(1);
152        if context.position > depth_budget {
153            RuleResult::stop_with_bound(f64::INFINITY, Reason::RuleReason)
154        } else {
155            RuleResult::continue_search()
156        }
157    }
158
159    fn priority(&self) -> u8 {
160        self.priority
161    }
162
163    fn description(&self) -> String {
164        "TopK rule".to_string()
165    }
166
167    fn state(&self) -> RuleState {
168        self.state
169    }
170
171    fn is_active(&self) -> bool {
172        self.state == RuleState::Active
173    }
174
175    fn activate(&mut self) {
176        self.state = RuleState::Active
177    }
178
179    fn deactivate(&mut self) {
180        self.state = RuleState::Disabled
181    }
182
183    fn relax(&mut self) {
184        if !self.is_active() {
185            return;
186        }
187
188        if self.is_relaxable() && self.budget >= self.limit {
189            self.deactivate();
190            return;
191        }
192        self.budget = self.increment.next();
193        if self.budget >= self.limit {
194            self.budget = self.limit;
195        }
196    }
197
198    fn is_relaxable(&self) -> bool {
199        self.relaxable
200    }
201
202    fn delay(&self) -> u8 {
203        self.delay
204    }
205
206    fn as_any(&self) -> &dyn std::any::Any {
207        self
208    }
209
210    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
211        self
212    }
213}