khive_brain_core/
section_state.rs

1//! Section posterior state — Thompson sampling and deterministic weight derivation.
2
3use std::collections::HashMap;
4
5use rand::Rng;
6use serde::{Deserialize, Serialize};
7
8use crate::brain_signal::BrainSignal;
9use crate::posterior::BetaPosterior;
10use crate::section_type::SectionType;
11use crate::signal::FeedbackSignal;
12
13pub const DEFAULT_ESS_CAP: f64 = 100.0;
14pub const DEFAULT_EXPLORATION_EPOCH: u64 = 50;
15pub const DEFAULT_TAU_0: f64 = 1.0;
16pub const DEFAULT_TAU_EXPLOIT: f64 = 0.1;
17pub const DEFAULT_SECTION_WEIGHT_FLOOR: f64 = 0.05;
18
19/// Per-profile section posterior state.
20pub struct SectionPosteriorState {
21    pub posteriors: HashMap<SectionType, BetaPosterior>,
22    pub priors: HashMap<SectionType, BetaPosterior>,
23    pub total_events: u64,
24    pub exploration_epoch: u64,
25}
26
27impl SectionPosteriorState {
28    pub fn new() -> Self {
29        let priors = Self::default_priors();
30        let posteriors = priors.clone();
31        Self {
32            posteriors,
33            priors,
34            total_events: 0,
35            exploration_epoch: DEFAULT_EXPLORATION_EPOCH,
36        }
37    }
38
39    pub fn from_priors(mut priors: HashMap<SectionType, BetaPosterior>) -> Self {
40        let neutral = BetaPosterior::new(2.0, 2.0);
41        for &st in SectionType::all() {
42            priors.entry(st).or_insert_with(|| neutral.clone());
43        }
44        let posteriors = priors.clone();
45        Self {
46            posteriors,
47            priors,
48            total_events: 0,
49            exploration_epoch: DEFAULT_EXPLORATION_EPOCH,
50        }
51    }
52
53    pub fn default_priors() -> HashMap<SectionType, BetaPosterior> {
54        let mut m = HashMap::new();
55        m.insert(SectionType::Overview, BetaPosterior::new(2.0, 2.0));
56        m.insert(SectionType::CoreModel, BetaPosterior::new(4.0, 2.0));
57        m.insert(
58            SectionType::BoundaryConditions,
59            BetaPosterior::new(2.0, 3.0),
60        );
61        m.insert(SectionType::Formalism, BetaPosterior::new(1.5, 4.0));
62        m.insert(
63            SectionType::OperationalGuidance,
64            BetaPosterior::new(6.0, 1.5),
65        );
66        m.insert(SectionType::Examples, BetaPosterior::new(5.0, 2.0));
67        m.insert(SectionType::FailureModes, BetaPosterior::new(3.0, 2.0));
68        m.insert(SectionType::ExpertLens, BetaPosterior::new(3.0, 2.0));
69        m.insert(SectionType::References, BetaPosterior::new(2.0, 2.0));
70        m.insert(SectionType::Other, BetaPosterior::new(2.0, 2.0));
71        m
72    }
73
74    pub fn to_snapshot(&self) -> SectionPosteriorSnapshot {
75        SectionPosteriorSnapshot {
76            posteriors: self.posteriors.clone(),
77            priors: self.priors.clone(),
78            total_events: self.total_events,
79            exploration_epoch: self.exploration_epoch,
80        }
81    }
82
83    pub fn from_snapshot(snapshot: SectionPosteriorSnapshot) -> Self {
84        Self {
85            posteriors: snapshot.posteriors,
86            priors: snapshot.priors,
87            total_events: snapshot.total_events,
88            exploration_epoch: snapshot.exploration_epoch,
89        }
90    }
91
92    pub fn weights(&self, rng: &mut impl Rng) -> HashMap<SectionType, f64> {
93        if self.exploration_epoch > 0 {
94            self.sample_weights(rng)
95        } else {
96            self.deterministic_weights()
97        }
98    }
99
100    pub fn sample_weights(&self, rng: &mut impl Rng) -> HashMap<SectionType, f64> {
101        let tau =
102            DEFAULT_TAU_0 * (self.exploration_epoch as f64 / DEFAULT_EXPLORATION_EPOCH as f64);
103        let tau = tau.max(1e-6);
104
105        let mut samples: HashMap<SectionType, f64> = HashMap::new();
106        for (&st, posterior) in &self.posteriors {
107            let theta =
108                sample_beta_gamma(posterior.alpha().max(1e-6), posterior.beta().max(1e-6), rng);
109            samples.insert(st, theta / tau);
110        }
111
112        let max_val = samples.values().cloned().fold(f64::NEG_INFINITY, f64::max);
113        let mut raw: HashMap<SectionType, f64> = HashMap::new();
114        for (&st, &logit) in &samples {
115            raw.insert(st, (logit - max_val).exp());
116        }
117
118        apply_floor_and_renorm(&mut raw, DEFAULT_SECTION_WEIGHT_FLOOR);
119        raw
120    }
121
122    pub fn deterministic_weights(&self) -> HashMap<SectionType, f64> {
123        let tau = DEFAULT_TAU_EXPLOIT;
124
125        let logits: HashMap<SectionType, f64> = self
126            .posteriors
127            .iter()
128            .map(|(&st, p)| (st, p.mean() / tau))
129            .collect();
130        let max_val = logits.values().cloned().fold(f64::NEG_INFINITY, f64::max);
131        let mut raw: HashMap<SectionType, f64> = HashMap::new();
132        for (&st, &logit) in &logits {
133            raw.insert(st, (logit - max_val).exp());
134        }
135
136        apply_floor_and_renorm(&mut raw, DEFAULT_SECTION_WEIGHT_FLOOR);
137        raw
138    }
139
140    pub fn reset_posteriors(&mut self) {
141        self.posteriors = self.priors.clone();
142    }
143
144    /// Apply a brain signal to update section posteriors in place.
145    /// Only `Feedback` events with `section_signals` affect section state.
146    pub fn apply_signal(&mut self, signal: &BrainSignal) {
147        if let BrainSignal::Feedback {
148            section_signals: Some(ref signals),
149            ..
150        } = signal
151        {
152            self.total_events += 1;
153
154            for (section_type, feedback_signal) in signals {
155                if let Some(posterior) = self.posteriors.get_mut(section_type) {
156                    match feedback_signal {
157                        FeedbackSignal::Useful => posterior.update_success(),
158                        FeedbackSignal::NotUseful => posterior.update_failure(),
159                        FeedbackSignal::Wrong => posterior.update_failure_weighted(2.0),
160                    }
161                    if let Some(prior) = self.priors.get(section_type) {
162                        if let Err(e) = posterior.apply_ess_cap(&prior.clone(), DEFAULT_ESS_CAP) {
163                            eprintln!(
164                                "[brain-core] apply_ess_cap failed for section {:?}: {e}",
165                                section_type
166                            );
167                        }
168                    }
169                }
170            }
171
172            if self.exploration_epoch > 0 {
173                self.exploration_epoch -= 1;
174            }
175        }
176    }
177}
178
179impl Default for SectionPosteriorState {
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185/// Derive section weights (Thompson or deterministic depending on epoch).
186pub fn derive_weights(
187    state: &SectionPosteriorState,
188    rng: &mut impl Rng,
189) -> HashMap<SectionType, f64> {
190    state.weights(rng)
191}
192
193/// Derive deterministic weights from posterior means.
194pub fn derive_deterministic_weights(state: &SectionPosteriorState) -> HashMap<SectionType, f64> {
195    state.deterministic_weights()
196}
197
198// ── Snapshot ────────────────────────────────────────────────────────────────
199
200#[derive(Debug, Clone, Serialize, Deserialize)]
201pub struct SectionPosteriorSnapshot {
202    pub posteriors: HashMap<SectionType, BetaPosterior>,
203    pub priors: HashMap<SectionType, BetaPosterior>,
204    pub total_events: u64,
205    pub exploration_epoch: u64,
206}
207
208// ── Sampling helpers ────────────────────────────────────────────────────────
209
210fn apply_floor_and_renorm(weights: &mut HashMap<SectionType, f64>, floor: f64) {
211    let sum: f64 = weights.values().sum();
212    if sum > 0.0 {
213        for v in weights.values_mut() {
214            *v /= sum;
215        }
216    }
217    for _ in 0..20 {
218        let (pinned_sum, n_free) = weights.values().fold((0.0f64, 0usize), |(ps, nf), &w| {
219            if w <= floor {
220                (ps + floor, nf)
221            } else {
222                (ps, nf + 1)
223            }
224        });
225        if n_free == 0 {
226            let total: f64 = weights.values().sum();
227            if total > 0.0 {
228                for v in weights.values_mut() {
229                    *v /= total;
230                }
231            }
232            break;
233        }
234        let free_mass = (1.0 - pinned_sum).max(0.0);
235        let free_sum: f64 = weights.values().filter(|&&w| w > floor).sum();
236        for v in weights.values_mut() {
237            if *v <= floor {
238                *v = floor;
239            } else if free_sum > 0.0 {
240                *v = (*v / free_sum) * free_mass;
241            }
242        }
243        if weights.values().all(|&w| w >= floor - 1e-12) {
244            break;
245        }
246    }
247}
248
249fn sample_beta_gamma(alpha: f64, beta: f64, rng: &mut impl Rng) -> f64 {
250    let x = sample_gamma_mt(alpha, rng);
251    let y = sample_gamma_mt(beta, rng);
252    let s = x + y;
253    if s <= 0.0 {
254        0.5
255    } else {
256        x / s
257    }
258}
259
260fn sample_gamma_mt(shape: f64, rng: &mut impl Rng) -> f64 {
261    if shape < 1.0 {
262        let g = sample_gamma_mt(shape + 1.0, rng);
263        let u: f64 = rng.gen();
264        return g * u.powf(1.0 / shape);
265    }
266    let d = shape - 1.0 / 3.0;
267    let c = 1.0 / (9.0 * d).sqrt();
268    loop {
269        let x: f64 = sample_standard_normal_bm(rng);
270        let v_raw = 1.0 + c * x;
271        if v_raw <= 0.0 {
272            continue;
273        }
274        let v = v_raw * v_raw * v_raw;
275        let u: f64 = rng.gen();
276        if u < 1.0 - 0.0331 * (x * x) * (x * x) {
277            return d * v;
278        }
279        if u.ln() < 0.5 * x * x + d * (1.0 - v + v.ln()) {
280            return d * v;
281        }
282    }
283}
284
285fn sample_standard_normal_bm(rng: &mut impl Rng) -> f64 {
286    let u1: f64 = rng.gen::<f64>().max(f64::EPSILON);
287    let u2: f64 = rng.gen();
288    (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn derive_weights_sums_to_one() {
297        let state = SectionPosteriorState::new();
298        let mut rng = rand::thread_rng();
299        for _ in 0..20 {
300            let weights = derive_weights(&state, &mut rng);
301            assert_eq!(weights.len(), SectionType::all().len());
302            let sum: f64 = weights.values().sum();
303            assert!(
304                (sum - 1.0).abs() < 1e-9,
305                "weights must sum to 1.0; got {sum}"
306            );
307        }
308    }
309
310    #[test]
311    fn deterministic_weights_sum_to_one() {
312        let state = SectionPosteriorState::new();
313        let weights = derive_deterministic_weights(&state);
314        assert_eq!(weights.len(), 10);
315        let sum: f64 = weights.values().sum();
316        assert!(
317            (sum - 1.0).abs() < 1e-9,
318            "deterministic weights must sum to 1.0; got {sum}"
319        );
320    }
321
322    #[test]
323    fn epoch_zero_uses_deterministic() {
324        let mut state = SectionPosteriorState::new();
325        state.exploration_epoch = 0;
326        let mut rng = rand::thread_rng();
327        let w1 = derive_weights(&state, &mut rng);
328        let w2 = derive_weights(&state, &mut rng);
329        for st in SectionType::all() {
330            assert!(
331                (w1[st] - w2[st]).abs() < 1e-12,
332                "epoch=0 should give deterministic weights for {:?}",
333                st
334            );
335        }
336    }
337
338    #[test]
339    fn higher_alpha_gets_higher_weight() {
340        let mut priors = SectionPosteriorState::default_priors();
341        priors.insert(
342            SectionType::OperationalGuidance,
343            BetaPosterior::new(20.0, 2.0),
344        );
345        let mut state = SectionPosteriorState::from_priors(priors);
346        state.exploration_epoch = 0;
347        let weights = derive_deterministic_weights(&state);
348        let og = weights[&SectionType::OperationalGuidance];
349        let form = weights[&SectionType::Formalism];
350        assert!(og > form, "og={og:.4} form={form:.4}");
351    }
352}
khive_brain_core/section_state.rs

khive_brain_core/
section_state.rs
