1use crate::activation::{ActivationConfig, EffectContext, PairStats};
4use crate::arrow_io::OutcomesRef;
5use crate::table::{ColGraph, ColumnVec};
6use ndarray::{Array1, Array2};
7use std::collections::hash_map::Entry;
8use std::collections::{HashMap, HashSet};
9
10#[inline]
11fn canonical_val_pair(a: i32, b: i32) -> (i32, i32) {
12 if a <= b {
13 (a, b)
14 } else {
15 (b, a)
16 }
17}
18
19#[derive(Clone, Debug)]
20pub struct PairAggregator {
21 pub target: String,
22 pub num_chunks: u64,
23 avg_outcome_sum: f32,
24 pub avg_count: u64,
25 k: i32,
26 col_map_int: HashMap<i32, String>,
27 col_map_str: HashMap<String, i32>,
28 val_map_int: HashMap<i32, String>,
29 pub val_map_str: HashMap<String, i32>,
30 col_graph_names: Vec<String>,
31 cols_dropped: HashSet<usize>,
32 pub cols: Vec<usize>,
33 col_names_ordered: Vec<String>,
34 vals_map: HashMap<(i32, i32), [f32; 2]>,
35 pub vals_map_avg: HashMap<(i32, i32), f32>,
36 pub avg_outcome: f32,
37 pub col_array: Vec<usize>,
38 tup_combos: HashMap<usize, (usize, usize)>,
39 col_to_tup: HashMap<usize, Vec<(usize, usize)>>,
40 m_divisor: f32,
41 combos_initialized: bool,
42 activation: ActivationConfig,
43}
44
45impl PairAggregator {
46 pub fn new() -> Self {
47 Self {
48 target: String::new(),
49 num_chunks: 0,
50 avg_outcome_sum: 0.0,
51 avg_count: 0,
52 k: 0,
53 col_map_int: HashMap::new(),
54 col_map_str: HashMap::new(),
55 val_map_int: HashMap::new(),
56 val_map_str: HashMap::new(),
57 col_graph_names: Vec::new(),
58 cols_dropped: HashSet::new(),
59 cols: Vec::new(),
60 col_names_ordered: Vec::new(),
61 vals_map: HashMap::new(),
62 vals_map_avg: HashMap::new(),
63 avg_outcome: 0.0,
64 col_array: Vec::new(),
65 tup_combos: HashMap::new(),
66 col_to_tup: HashMap::new(),
67 m_divisor: 1.0,
68 combos_initialized: false,
69 activation: ActivationConfig::default(),
70 }
71 }
72
73 pub fn with_activation(activation: ActivationConfig) -> Self {
74 Self {
75 activation,
76 ..Self::new()
77 }
78 }
79
80 pub fn vals_map_len(&self) -> usize {
81 self.vals_map.len()
82 }
83
84 fn combo_pair(&self, combo_id: usize) -> Result<(usize, usize), String> {
85 self.tup_combos
86 .get(&combo_id)
87 .copied()
88 .ok_or_else(|| format!("internal: unknown combo id {combo_id}"))
89 }
90
91 fn mapped_col<'a>(
92 &self,
93 x_mapped: &'a HashMap<usize, Array1<i32>>,
94 col_idx: usize,
95 ) -> Result<&'a Array1<i32>, String> {
96 x_mapped
97 .get(&col_idx)
98 .ok_or_else(|| format!("internal: missing mapped column {col_idx}"))
99 }
100
101 pub fn initialize_inputs(
102 &mut self,
103 col_info: &ColGraph,
104 target: &str,
105 column_order: &[String],
106 ) -> Result<(), String> {
107 self.target = target.to_string();
108 self.num_chunks = 0;
109 self.avg_outcome_sum = 0.0;
110 self.avg_count = 0;
111 self.k = 0;
112 self.col_map_int.clear();
113 self.col_map_str.clear();
114 self.val_map_int.clear();
115 self.val_map_str.clear();
116 self.col_graph_names = col_info.names.clone();
117 self.cols_dropped = col_info.dropped.clone();
118 self.cols = col_info.active_indices();
119 self.col_names_ordered = column_order.to_vec();
120 self.combos_initialized = false;
121 self.col_array.clear();
122 self.tup_combos.clear();
123 self.col_to_tup.clear();
124
125 for name in column_order {
126 let s = name.clone();
127 if !self.col_map_str.contains_key(&s) {
128 let id = self.k;
129 self.col_map_int.insert(id, s.clone());
130 self.col_map_str.insert(s, id);
131 self.k += 1;
132 }
133 }
134 Ok(())
135 }
136
137 pub fn make_col_combos(&mut self) {
138 if self.combos_initialized {
139 return;
140 }
141 self.vals_map.clear();
142 let mut pairs: Vec<(usize, usize)> = Vec::new();
143 let idxs = &self.cols;
144 for i in 0..idxs.len() {
145 for j in (i + 1)..idxs.len() {
146 pairs.push((idxs[i], idxs[j]));
147 }
148 }
149 for (k, &(a, b)) in pairs.iter().enumerate() {
150 self.col_array.push(k);
151 self.tup_combos.insert(k, (a, b));
152 }
153 let mut col_to_tup: HashMap<usize, Vec<(usize, usize)>> = HashMap::new();
154 for (&_c, &(a, b)) in self.tup_combos.iter() {
155 let srt = if a < b { (a, b) } else { (b, a) };
156 col_to_tup.entry(a).or_default().push(srt);
157 col_to_tup.entry(b).or_default().push(srt);
158 }
159 self.col_to_tup = col_to_tup
160 .into_iter()
161 .map(|(c, mut v)| {
162 v.sort();
163 v.dedup();
164 (c, v)
165 })
166 .collect();
167
168 self.m_divisor = (self.col_to_tup.len() as f32 - 1.0).max(1.0);
169 self.combos_initialized = true;
170 }
171
172 fn ensure_sentinel_values(&mut self) {
173 if self.num_chunks == 0 {
174 if !self.val_map_str.contains_key("no data") {
175 let id = self.k;
176 self.val_map_int.insert(id, "no data".into());
177 self.val_map_str.insert("no data".into(), id);
178 self.k += 1;
179 }
180 if !self.val_map_str.contains_key("None") {
181 let id = self.k;
182 self.val_map_int.insert(id, "None".into());
183 self.val_map_str.insert("None".into(), id);
184 self.k += 1;
185 }
186 }
187 }
188
189 fn column_vec_to_labels(col: &ColumnVec) -> Result<Vec<String>, String> {
190 match col {
191 ColumnVec::Utf8(v) => Ok(v
192 .iter()
193 .map(|s| {
194 let t = s.as_str();
195 if t.is_empty() {
196 "no data".to_string()
197 } else {
198 t.to_string()
199 }
200 })
201 .collect()),
202 ColumnVec::F32(v) => Ok(v
203 .iter()
204 .map(|&x| {
205 if x.is_finite() {
206 x.to_string()
207 } else {
208 "0".to_string()
209 }
210 })
211 .collect()),
212 ColumnVec::F32Array(v) => Ok(v
213 .iter()
214 .map(|&x| {
215 if x.is_finite() {
216 x.to_string()
217 } else {
218 "0".to_string()
219 }
220 })
221 .collect()),
222 }
223 }
224
225 fn register_column_values(&mut self, labels: &[String]) -> Result<(), String> {
226 self.ensure_sentinel_values();
227 let mut uniq: Vec<&String> = labels.iter().collect();
228 uniq.sort();
229 uniq.dedup();
230 for val in uniq {
231 if let Entry::Vacant(e) = self.val_map_str.entry(val.clone()) {
232 let id = self.k;
233 self.val_map_int.insert(id, val.clone());
234 e.insert(id);
235 self.k += 1;
236 }
237 }
238 Ok(())
239 }
240
241 fn column_to_ids(&self, labels: &[String]) -> Result<Array1<i32>, String> {
242 let mut v = Vec::with_capacity(labels.len());
243 for val in labels {
244 let id = self
245 .val_map_str
246 .get(val)
247 .copied()
248 .ok_or_else(|| format!("unknown val {val}"))?;
249 v.push(id);
250 }
251 Ok(Array1::from(v))
252 }
253
254 fn x_processed_to_mapped(
255 &mut self,
256 x: &HashMap<String, ColumnVec>,
257 ) -> Result<HashMap<usize, Array1<i32>>, String> {
258 let n = x
259 .get(&self.col_graph_names[0])
260 .map(|c| c.len())
261 .ok_or_else(|| "missing first col".to_string())?;
262 let col_indices: Vec<usize> = self.cols.clone();
263 let mut out: HashMap<usize, Array1<i32>> = HashMap::new();
264 for col_idx in col_indices {
265 let name = self.col_graph_names[col_idx].clone();
266 let col = x.get(&name).ok_or_else(|| format!("missing col {name}"))?;
267 let labels = Self::column_vec_to_labels(col)?;
268 if labels.len() != n {
269 return Err("column length mismatch in x_processed_to_mapped".into());
270 }
271 self.register_column_values(&labels)?;
272 let ids = self.column_to_ids(&labels)?;
273 out.insert(col_idx, ids);
274 }
275 Ok(out)
276 }
277
278 pub fn vals_map_updating(
279 &mut self,
280 x_processed: &HashMap<String, ColumnVec>,
281 outcomes: &OutcomesRef<'_>,
282 ) -> Result<(), String> {
283 let x_mapped = self.x_processed_to_mapped(x_processed)?;
284
285 self.avg_outcome_sum += outcomes.sum();
286 self.avg_count += outcomes.len() as u64;
287
288 let n = outcomes.len();
289 let one_percent = ((n as f32) * 0.01).floor() as usize;
290
291 for &c in &self.col_array {
292 let (c1, c2) = self.combo_pair(c)?;
293 let a = self.mapped_col(&x_mapped, c1)?;
294 let b = self.mapped_col(&x_mapped, c2)?;
295
296 let mut local: HashMap<(i32, i32), PairStats> = HashMap::new();
297 for i in 0..n {
298 let key = canonical_val_pair(a[i], b[i]);
299 let out_i = {
300 let x = outcomes.get(i);
301 if x.is_nan() {
302 0.0
303 } else {
304 x
305 }
306 };
307 let entry = local.entry(key).or_insert(PairStats {
308 sum: 0.0,
309 count: 0.0,
310 });
311 entry.sum += out_i;
312 entry.count += 1.0;
313 }
314
315 for (key, stats) in local {
316 if stats.count as usize > one_percent {
317 let entry = self.vals_map.entry(key).or_insert([0.0, 0.0]);
318 entry[0] += stats.sum;
319 entry[1] += stats.count;
320 } else {
321 self.vals_map.entry(key).or_insert([0.0, 0.0]);
322 }
323 }
324 }
325
326 self.num_chunks += 1;
327 Ok(())
328 }
329
330 pub fn finish_map(&mut self) {
331 self.vals_map_avg.clear();
332 for (&key, &arr) in &self.vals_map {
333 let weight = self.activation.kg_pair.activate(PairStats {
334 sum: arr[0],
335 count: arr[1],
336 });
337 self.vals_map_avg.insert(key, weight);
338 }
339
340 if self.avg_count > 0 {
341 self.avg_outcome = self.avg_outcome_sum / self.avg_count as f32;
342 } else {
343 self.avg_outcome = 0.0;
344 }
345 }
346
347 fn make_cvto_inner(
348 &mut self,
349 x_processed: &HashMap<String, ColumnVec>,
350 ) -> Result<Array2<f32>, String> {
351 let x_mapped = self.x_processed_to_mapped(x_processed)?;
352 let n = x_mapped
353 .get(&self.cols[0])
354 .map(|a| a.len())
355 .ok_or_else(|| "no active columns".to_string())?;
356 let m = self.col_array.len();
357 let mut col_vals = Array2::<f32>::zeros((n, m));
358
359 for (mi, &combo_id) in self.col_array.iter().enumerate() {
360 let (c1, c2) = self.combo_pair(combo_id)?;
361 let a = self.mapped_col(&x_mapped, c1)?;
362 let b = self.mapped_col(&x_mapped, c2)?;
363 for i in 0..n {
364 let tup = canonical_val_pair(a[i], b[i]);
365 let v = *self.vals_map_avg.get(&tup).unwrap_or(&0.0);
366 col_vals[[i, mi]] = v;
367 }
368 }
369 Ok(col_vals)
370 }
371
372 pub fn use_map(
373 &mut self,
374 mut x_processed: HashMap<String, ColumnVec>,
375 y: Vec<f32>,
376 outcomes: Vec<f32>,
377 ) -> Result<HashMap<String, ColumnVec>, String> {
378 let col_vals_outcomes = self.make_cvto_inner(&x_processed)?;
379 let n = col_vals_outcomes.nrows();
380
381 let mut col_combined: HashMap<usize, Array1<f32>> = HashMap::new();
382 for &combo_id in &self.col_array {
383 let (c1, c2) = self.combo_pair(combo_id)?;
384 let col_view = col_vals_outcomes.column(combo_id);
385 for &col_idx in &[c1, c2] {
386 col_combined
387 .entry(col_idx)
388 .or_insert_with(|| Array1::zeros(n))
389 .scaled_add(1.0, &col_view);
390 }
391 }
392 let m = self.m_divisor.max(1.0);
393 for v in col_combined.values_mut() {
394 *v /= m;
395 }
396
397 let mut nnm: HashMap<String, ColumnVec> = HashMap::new();
398 for name in &self.col_graph_names {
399 let colvec = x_processed
400 .remove(name)
401 .unwrap_or_else(|| ColumnVec::Utf8(vec!["no data".into(); n]));
402 nnm.insert(name.clone(), colvec);
403 }
404
405 let ctx = EffectContext {
406 global_mean_outcome: self.avg_outcome,
407 };
408 for c_idx in 0..self.col_graph_names.len() {
409 if let Some(arr) = col_combined.get(&c_idx) {
410 let col_name = &self.col_graph_names[c_idx];
411 let diffs = Array1::from(
412 arr.iter()
413 .map(|x| self.activation.effect.activate(*x, &ctx))
414 .collect::<Vec<f32>>(),
415 );
416 nnm.insert(format!("{col_name}_effect"), ColumnVec::F32Array(diffs));
417 }
418 }
419
420 nnm.insert("Actuals".into(), ColumnVec::F32Array(Array1::from(y)));
421 nnm.insert(
422 "outcomes_effect".into(),
423 ColumnVec::F32Array(Array1::from(outcomes)),
424 );
425 Ok(nnm)
426 }
427}
428
429impl Default for PairAggregator {
430 fn default() -> Self {
431 Self::new()
432 }
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438 use std::collections::HashMap;
439
440 fn two_col_aggregator() -> PairAggregator {
441 let cg = ColGraph {
442 names: vec!["a".into(), "b".into()],
443 dropped: HashSet::new(),
444 };
445 let mut agg = PairAggregator::new();
446 agg.initialize_inputs(&cg, "target", &["a".into(), "b".into()])
447 .unwrap();
448 agg.make_col_combos();
449 agg
450 }
451
452 fn oriented_pair_data() -> HashMap<String, ColumnVec> {
453 let mut x = HashMap::new();
454 x.insert("a".into(), ColumnVec::Utf8(vec!["5".into(), "7".into()]));
455 x.insert("b".into(), ColumnVec::Utf8(vec!["7".into(), "5".into()]));
456 x
457 }
458
459 fn outcomes_slice(v: &[f32]) -> crate::arrow_io::OutcomesRef<'_> {
460 crate::arrow_io::OutcomesRef::Slice(v)
461 }
462
463 #[test]
464 fn canonical_val_pair_commutes() {
465 assert_eq!(canonical_val_pair(5, 7), canonical_val_pair(7, 5));
466 assert_eq!(canonical_val_pair(5, 7), (5, 7));
467 }
468
469 #[test]
470 fn vals_map_merges_orientations() {
471 let mut agg = two_col_aggregator();
472 let x = oriented_pair_data();
473 agg.vals_map_updating(&x, &outcomes_slice(&[1.0, 2.0]))
474 .unwrap();
475
476 let id5 = agg.val_map_str["5"];
477 let id7 = agg.val_map_str["7"];
478 let key = canonical_val_pair(id5, id7);
479
480 assert_eq!(agg.vals_map.len(), 1);
481 let arr = agg.vals_map[&key];
482 assert!((arr[0] - 3.0).abs() < 1e-5);
483 assert!((arr[1] - 2.0).abs() < 1e-5);
484 }
485
486 #[test]
487 fn finish_map_no_duplicate_keys() {
488 let mut agg = two_col_aggregator();
489 let x = oriented_pair_data();
490 agg.vals_map_updating(&x, &outcomes_slice(&[1.0, 2.0]))
491 .unwrap();
492 agg.finish_map();
493 assert_eq!(agg.vals_map_avg.len(), agg.vals_map.len());
494 }
495
496 #[test]
497 fn lookup_symmetric() {
498 let mut agg = two_col_aggregator();
499 let x = oriented_pair_data();
500 agg.vals_map_updating(&x, &outcomes_slice(&[1.0, 2.0]))
501 .unwrap();
502 agg.finish_map();
503
504 let col_vals = agg.make_cvto_inner(&x).unwrap();
505 assert_eq!(col_vals.nrows(), 2);
506 assert!((col_vals[[0, 0]] - col_vals[[1, 0]]).abs() < 1e-5);
507 }
508}