1use crate::sparse_data_visitors::*;
2use crate::sparse_io::*;
3use crate::sparse_io_vector::*;
4
5use indicatif::ParallelProgressIterator;
6use legume_numeric::matrix::sparse_stat::{SparseColumnRunningStatistics, SparseRunningStatistics};
7use legume_numeric::matrix::traits::RunningStatOps;
8use legume_numeric::matrix::utils::partition_by_membership;
9use log::warn;
10use rayon::prelude::*;
11use std::sync::{Arc, Mutex};
12
13#[allow(dead_code)]
16pub mod axis;
17pub use axis::Axis;
18
19use rustc_hash::FxHashMap as HashMap;
20
21#[derive(Clone)]
22pub struct SqueezeCutoffs {
23 pub row: usize,
24 pub column: usize,
25}
26
27pub fn below_nnz_cutoff(nnz: f32, cutoff: usize) -> bool {
31 (nnz as usize) < cutoff
32}
33
34pub fn squeeze_by_nnz(
42 data: &dyn SparseIo<IndexIter = Vec<usize>>,
43 cutoffs: SqueezeCutoffs,
44 block_size: Option<usize>,
45 preload: bool,
46) -> anyhow::Result<()> {
47 anyhow::ensure!(
49 data.num_rows().is_some() && data.num_columns().is_some(),
50 "missing shape information"
51 );
52 let (col_stat, row_stat) = collect_column_row_stat(data, block_size)?;
53
54 fn nnz_index(nnz: &[f32], cutoff: usize) -> Option<Vec<usize>> {
55 let ret: Vec<usize> = nnz
56 .iter()
57 .enumerate()
58 .filter(|&(_, &x)| !below_nnz_cutoff(x, cutoff))
59 .map(|(i, _)| i)
60 .collect();
61
62 (!ret.is_empty()).then_some(ret)
63 }
64
65 let row_nnz_vec = row_stat.count_positives();
66 let col_nnz_vec = col_stat.count_positives();
67 let row_idx = nnz_index(&row_nnz_vec, cutoffs.row);
68 let col_idx = nnz_index(&col_nnz_vec, cutoffs.column);
69
70 if row_idx.is_none() {
71 warn!(
72 "No rows can be kept with this cutoff {}!\n\
73 \n\
74 We will stop squeezing on the rows.\n\
75 \n",
76 cutoffs.row
77 );
78 }
79
80 if col_idx.is_none() {
81 warn!(
82 "No columns can be kept with this cutoff {}!\n\
83 \n\
84 We will stop squeezing on the columns.\n\
85 \n",
86 cutoffs.column
87 );
88 }
89
90 let row_idx = row_idx.filter(|kept| kept.len() < row_nnz_vec.len());
93 let col_idx = col_idx.filter(|kept| kept.len() < col_nnz_vec.len());
94 if row_idx.is_none() && col_idx.is_none() {
95 return Ok(());
96 }
97
98 let file = data.get_backend_file_name();
99 let backend = data.backend_type();
100
101 let mut data = open_sparse_matrix(file, &backend)?;
102 if preload {
103 data.preload_columns()?;
104 }
105 data.subset_columns_rows(col_idx.as_ref(), row_idx.as_ref())
106}
107
108pub fn collect_row_stat_across_vec(
112 data: &SparseIoVec,
113 block_size: Option<usize>,
114) -> anyhow::Result<SparseRunningStatistics<f32>> {
115 let mut row_stat = SparseRunningStatistics::new(data.num_rows());
116 data.visit_columns_by_block(
117 &row_stat_vec_visitor,
118 &EmptyArgs {},
119 &mut row_stat,
120 block_size,
121 )?;
122 Ok(row_stat)
123}
124
125#[allow(clippy::type_complexity)]
130pub fn collect_stratified_row_stat_across_vec(
131 data: &SparseIoVec,
132 column_membership: &HashMap<Box<str>, Box<str>>,
133 block_size: Option<usize>,
134) -> anyhow::Result<(Vec<Box<str>>, Vec<SparseRunningStatistics<f32>>)> {
135 let column_names = data.column_names()?;
136 let default = "".to_string().into_boxed_str();
137 let membership = column_names
138 .into_iter()
139 .map(|k| column_membership.get(&k).unwrap_or(&default).clone())
140 .collect::<Vec<_>>();
141
142 let partitions = partition_by_membership(&membership, None);
143 let mut group_names = Vec::with_capacity(partitions.len());
144 let mut group_stats = Vec::with_capacity(partitions.len());
145 let num_features = data.num_rows();
146
147 for (k, cols) in partitions {
148 let jobs = create_jobs(cols.len(), num_features, block_size);
149 let mut row_stat = SparseRunningStatistics::new(data.num_rows());
150 let arc_stat = Arc::new(Mutex::new(&mut row_stat));
151
152 jobs.par_iter()
153 .progress_with(styled_progress_bar(jobs.len() as u64, "blocks"))
154 .for_each(|&(lb, ub)| {
155 let cols_sub = cols[lb..ub].iter().cloned();
156 let csc = data
157 .read_columns_csc(cols_sub)
158 .expect("failed to read data");
159 let mut stat = arc_stat.lock().expect("failed to lock row_stat");
160 stat.add_csc(&csc);
161 });
162
163 group_names.push(k);
164 group_stats.push(row_stat);
165 }
166
167 Ok((group_names, group_stats))
168}
169
170pub fn collect_row_stat(
174 data: &dyn SparseIo<IndexIter = Vec<usize>>,
175 block_size: Option<usize>,
176) -> anyhow::Result<SparseRunningStatistics<f32>> {
177 let nrows = data.num_rows().unwrap_or(0);
178 fold_column_blocks(
179 data,
180 block_size,
181 || SparseRunningStatistics::new(nrows),
182 |stat, csc, _| stat.add_csc(csc),
183 |stat, other| stat.merge(&other),
184 )
185}
186
187pub(crate) fn blocks_per_leaf(n_blocks: usize) -> usize {
195 const LEAVES: usize = 256;
196 n_blocks.div_ceil(LEAVES).max(1)
197}
198
199pub(crate) fn fold_in_fixed_order<T, A, I, F, M>(
209 items: &[T],
210 leaf: usize,
211 init: &I,
212 fold: &F,
213 merge: &M,
214) -> anyhow::Result<A>
215where
216 T: Sync,
217 A: Send,
218 I: Fn() -> A + Sync,
219 F: Fn(&mut A, &T) -> anyhow::Result<()> + Sync,
220 M: Fn(&mut A, A) + Sync,
221{
222 if items.len() <= leaf.max(1) {
223 let mut acc = init();
224 for item in items {
225 fold(&mut acc, item)?;
226 }
227 return Ok(acc);
228 }
229 let (left, right) = items.split_at(items.len() / 2);
230 let (left, right) = rayon::join(
231 || fold_in_fixed_order(left, leaf, init, fold, merge),
232 || fold_in_fixed_order(right, leaf, init, fold, merge),
233 );
234 let mut acc = left?;
235 merge(&mut acc, right?);
236 Ok(acc)
237}
238
239fn fold_column_blocks<A, I, F, M>(
249 data: &dyn SparseIo<IndexIter = Vec<usize>>,
250 block_size: Option<usize>,
251 init: I,
252 add: F,
253 merge: M,
254) -> anyhow::Result<A>
255where
256 A: Send,
257 I: Fn() -> A + Sync,
258 F: Fn(&mut A, &CscMatrix<f32>, usize) + Sync,
259 M: Fn(&mut A, A) + Sync,
260{
261 let ncols = data.num_columns().unwrap_or(0);
262 let nrows = data.num_rows().unwrap_or(0);
263 let jobs = create_jobs(ncols, nrows, block_size);
264 let prog_bar = styled_progress_bar(jobs.len() as u64, "blocks");
265 fold_in_fixed_order(
266 &jobs,
267 blocks_per_leaf(jobs.len()),
268 &init,
269 &|acc, &(lb, ub)| {
270 let csc = data.read_columns_csc((lb..ub).collect())?;
271 add(acc, &csc, lb);
272 prog_bar.inc(1);
273 Ok(())
274 },
275 &merge,
276 )
277}
278
279pub fn masked_column_stat(
283 ncols: usize,
284 nrows: usize,
285 select_rows: &[usize],
286) -> (Vec<bool>, SparseColumnRunningStatistics<f32>) {
287 let mut mask = vec![false; nrows];
288 for &r in select_rows {
289 if r < nrows {
290 mask[r] = true;
291 }
292 }
293 let selected = mask.iter().filter(|x| **x).count();
294 (mask, SparseColumnRunningStatistics::new(ncols, selected))
295}
296
297pub fn collect_column_stat_across_vec(
302 data: &SparseIoVec,
303 select_rows: Option<&[usize]>,
304 block_size: Option<usize>,
305) -> anyhow::Result<SparseColumnRunningStatistics<f32>> {
306 let ncols = data.num_columns();
307 let nrows_total = data.num_rows();
308
309 let (row_mask, mut col_stat) = match select_rows {
310 Some(sel) => {
311 let (mask, stat) = masked_column_stat(ncols, nrows_total, sel);
312 (Some(mask), stat)
313 }
314 None => (None, SparseColumnRunningStatistics::new(ncols, nrows_total)),
315 };
316 data.visit_columns_by_block(&col_stat_visitor, &row_mask, &mut col_stat, block_size)?;
317 Ok(col_stat)
318}
319
320pub fn collect_column_row_stat(
325 data: &dyn SparseIo<IndexIter = Vec<usize>>,
326 block_size: Option<usize>,
327) -> anyhow::Result<(
328 SparseColumnRunningStatistics<f32>,
329 SparseRunningStatistics<f32>,
330)> {
331 let ncols = data.num_columns().unwrap_or(0);
332 let nrows = data.num_rows().unwrap_or(0);
333 let col_stat = Mutex::new(SparseColumnRunningStatistics::<f32>::new(ncols, nrows));
336 let row_stat = fold_column_blocks(
337 data,
338 block_size,
339 || SparseRunningStatistics::<f32>::new(nrows),
340 |row_stat, csc, lb| {
341 col_stat
342 .lock()
343 .expect("failed to lock the column statistics")
344 .add_csc(csc, lb);
345 row_stat.add_csc(csc);
346 },
347 |row_stat, other| row_stat.merge(&other),
348 )?;
349 let col_stat = col_stat
350 .into_inner()
351 .expect("failed to unlock the column statistics");
352 Ok((col_stat, row_stat))
353}
354
355pub fn collect_column_stat(
359 data: &dyn SparseIo<IndexIter = Vec<usize>>,
360 block_size: Option<usize>,
361) -> anyhow::Result<SparseColumnRunningStatistics<f32>> {
362 let ncols = data.num_columns().unwrap_or(0);
363 let nrows = data.num_rows().unwrap_or(0);
364 let col_stat = Mutex::new(SparseColumnRunningStatistics::<f32>::new(ncols, nrows));
367 fold_column_blocks(
368 data,
369 block_size,
370 || (),
371 |_, csc, lb| {
372 col_stat
373 .lock()
374 .expect("failed to lock the column statistics")
375 .add_csc(csc, lb)
376 },
377 |_, _| (),
378 )?;
379 Ok(col_stat
380 .into_inner()
381 .expect("failed to unlock the column statistics"))
382}
383
384struct EmptyArgs {}
385
386fn row_stat_vec_visitor(
387 job: (usize, usize),
388 data: &SparseIoVec,
389 _: &EmptyArgs,
390 arc_stat: Arc<Mutex<&mut SparseRunningStatistics<f32>>>,
391) -> anyhow::Result<()> {
392 let (lb, ub) = job;
393 let csc = data.read_columns_csc(lb..ub)?;
394
395 let mut stat = arc_stat.lock().expect("failed to lock row_stat");
396 stat.add_csc(&csc);
397 Ok(())
398}
399
400fn col_stat_visitor(
401 job: (usize, usize),
402 data: &SparseIoVec,
403 row_mask: &Option<Vec<bool>>,
404 arc_stat: Arc<Mutex<&mut SparseColumnRunningStatistics<f32>>>,
405) -> anyhow::Result<()> {
406 let (lb, ub) = job;
407 let csc = data.read_columns_csc(lb..ub)?;
408 let mut stat = arc_stat.lock().expect("failed to lock col_stat");
409 match row_mask {
410 Some(mask) => stat.add_csc_masked(&csc, lb, mask),
411 None => stat.add_csc(&csc, lb),
412 }
413 Ok(())
414}
415
416const TROUGH_BIN: f64 = 0.05;
423const TROUGH_SMOOTH_BINS: f64 = 3.0;
425const TROUGH_MAX_DEPTH: f64 = 0.25;
427const TROUGH_MIN_PEAK_RATIO: f64 = 10.0;
430const TROUGH_MIN_SIDE: usize = 20;
432
433pub fn suggest_nnz_cutoff(nnz: &[f32]) -> Option<usize> {
450 let n = nnz.len();
451 let min_side = TROUGH_MIN_SIDE.max(n / 1000);
452 if n < 2 * min_side {
453 return None;
454 }
455
456 let mut vals: Vec<u64> = nnz.iter().map(|&x| x.max(0.0).round() as u64).collect();
457 vals.sort_unstable();
458 let (vmin, vmax) = (vals[0], vals[n - 1]);
459 if vmin == vmax {
460 return None;
461 }
462
463 let edge = |v: u64, half: f64| (v as f64 + 1.0 + half).ln();
465 let lo = edge(vmin, -0.5);
466 let nbins = ((edge(vmax, 0.5) - lo) / TROUGH_BIN).ceil() as usize;
467 let mut hist = vec![0.0_f64; nbins];
468 let mut i = 0;
469 while i < n {
470 let v = vals[i];
471 let mut j = i;
472 while j < n && vals[j] == v {
473 j += 1;
474 }
475 let a = (edge(v, -0.5) - lo) / TROUGH_BIN;
476 let b = (edge(v, 0.5) - lo) / TROUGH_BIN;
477 let mass = (j - i) as f64 / (b - a);
478 for (k, h) in hist
479 .iter_mut()
480 .enumerate()
481 .take((b.ceil() as usize).min(nbins))
482 .skip(a.floor() as usize)
483 {
484 let overlap = b.min(k as f64 + 1.0) - a.max(k as f64);
485 if overlap > 0.0 {
486 *h += mass * overlap;
487 }
488 }
489 i = j;
490 }
491
492 let radius = (3.0 * TROUGH_SMOOTH_BINS).ceil() as isize;
494 let kernel: Vec<f64> = (-radius..=radius)
495 .map(|d| (-0.5 * (d as f64 / TROUGH_SMOOTH_BINS).powi(2)).exp())
496 .collect();
497 let smooth: Vec<f64> = (0..nbins as isize)
498 .map(|k| {
499 (-radius..=radius)
500 .filter_map(|d| {
501 let t = k + d;
502 (0..nbins as isize)
503 .contains(&t)
504 .then(|| hist[t as usize] * kernel[(d + radius) as usize])
505 })
506 .sum()
507 })
508 .collect();
509
510 let mut below = vec![0.0_f64; nbins + 1];
512 for k in 0..nbins {
513 below[k + 1] = below[k] + hist[k];
514 }
515 let mut left_peak = vec![(0.0_f64, 0usize); nbins];
516 for k in 0..nbins {
517 let prev = if k > 0 { left_peak[k - 1] } else { (-1.0, 0) };
518 left_peak[k] = if smooth[k] > prev.0 {
519 (smooth[k], k)
520 } else {
521 prev
522 };
523 }
524 let mut right_peak = vec![(0.0_f64, 0usize); nbins];
525 for k in (0..nbins).rev() {
526 let next = if k + 1 < nbins {
527 right_peak[k + 1]
528 } else {
529 (-1.0, k)
530 };
531 right_peak[k] = if smooth[k] >= next.0 {
532 (smooth[k], k)
533 } else {
534 next
535 };
536 }
537
538 let min_bins_apart = TROUGH_MIN_PEAK_RATIO.ln() / TROUGH_BIN;
539 let total = below[nbins];
540 let mut best: Option<(f64, usize, usize)> = None; for t in 1..nbins.saturating_sub(1) {
542 let (l_mass, r_mass) = (below[t], total - below[t + 1]);
543 if l_mass < min_side as f64 || r_mass < min_side as f64 {
544 continue;
545 }
546 let ((l_h, l_k), (r_h, r_k)) = (left_peak[t - 1], right_peak[t + 1]);
547 if ((r_k - l_k) as f64) < min_bins_apart || l_h <= 0.0 || r_h <= 0.0 {
548 continue;
549 }
550 let depth = smooth[t] / l_h.min(r_h);
551 match best {
552 Some((d, first, last)) if depth == d && last + 1 == t => best = Some((d, first, t)),
553 Some((d, _, _)) if depth >= d => {}
554 _ => best = Some((depth, t, t)),
555 }
556 }
557
558 let Some((depth, first, last)) = best else {
559 log::info!("nnz cell-calling: no two peaks a decade apart → unimodal, no cutoff");
560 return None;
561 };
562 let x = lo + ((first + last) as f64 / 2.0 + 0.5) * TROUGH_BIN;
564 let cutoff = (x.exp() - 1.0).ceil().max(1.0) as usize;
565 let favors_cut = depth <= TROUGH_MAX_DEPTH;
566 log::info!(
567 "nnz cell-calling: deepest trough at nnz {} (depth {:.3}) → {}",
568 cutoff,
569 depth,
570 if favors_cut {
571 format!("bimodal, cutoff at nnz {cutoff}")
572 } else {
573 "unimodal, no cutoff".to_string()
574 }
575 );
576 favors_cut.then_some(cutoff)
577}
578
579pub fn pct(part: usize, total: usize) -> f64 {
581 if total > 0 {
582 100.0 * part as f64 / total as f64
583 } else {
584 0.0
585 }
586}
587
588pub fn median_of_sorted(sorted: &[f32]) -> f32 {
590 let n = sorted.len();
591 match n {
592 0 => 0.0,
593 _ if n.is_multiple_of(2) => (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0,
594 _ => sorted[n / 2],
595 }
596}
597
598struct HistBin {
601 val_min: f32,
602 val_max: f32,
603 count: usize,
604 is_cutoff: bool,
605}
606
607const PRINTED_BINS: usize = 40;
609
610fn create_log_histogram(values: &[f32], cutoff: usize) -> Vec<HistBin> {
615 let axis = Axis::fit(values, true);
618 let axis = if cutoff > 0 {
619 axis.including(cutoff as f64)
620 } else {
621 axis
622 };
623 let mut bins: Vec<(usize, f32, f32)> =
624 vec![(0, f32::INFINITY, f32::NEG_INFINITY); PRINTED_BINS];
625 for &val in values.iter().filter(|v| v.is_finite()) {
626 let b = &mut bins[axis.bin(val as f64, PRINTED_BINS)];
627 b.0 += 1;
628 b.1 = b.1.min(val);
629 b.2 = b.2.max(val);
630 }
631 let cutoff_bin = (cutoff > 0)
635 .then(|| {
636 let at = axis.bin(cutoff as f64, PRINTED_BINS);
637 (at..PRINTED_BINS).find(|&k| bins[k].0 > 0)
638 })
639 .flatten();
640 bins.into_iter()
641 .enumerate()
642 .filter(|(_, b)| b.0 > 0)
643 .map(|(k, (count, val_min, val_max))| HistBin {
644 val_min,
645 val_max,
646 count,
647 is_cutoff: Some(k) == cutoff_bin,
648 })
649 .collect()
650}
651
652pub fn fmt_stat(v: f32, decimals: usize) -> String {
656 if v.fract() == 0.0 {
657 (v as i64).to_string()
658 } else {
659 format!("{:.*}", decimals, v)
660 }
661}
662
663fn fmt_sig(v: f32) -> String {
666 if v.fract() == 0.0 {
667 return (v as i64).to_string();
668 }
669 let places = (2 - (v.abs() as f64).log10().floor() as i32).max(0) as usize;
670 format!("{:.*}", places, v)
671}
672
673pub fn print_nnz_summary(
681 label: &str,
682 metric: &str,
683 values: &[f32],
684 cutoff: usize,
685 suggested: Option<usize>,
686) {
687 let _ = write_nnz_summary(
688 &mut std::io::stdout().lock(),
689 label,
690 metric,
691 values,
692 cutoff,
693 suggested,
694 );
695}
696
697pub fn eprint_nnz_summary(
701 label: &str,
702 metric: &str,
703 values: &[f32],
704 cutoff: usize,
705 suggested: Option<usize>,
706) {
707 let _ = write_nnz_summary(
708 &mut std::io::stderr().lock(),
709 label,
710 metric,
711 values,
712 cutoff,
713 suggested,
714 );
715}
716
717pub fn write_nnz_summary(
719 out: &mut impl std::io::Write,
720 label: &str,
721 metric: &str,
722 values: &[f32],
723 cutoff: usize,
724 suggested: Option<usize>,
725) -> std::io::Result<()> {
726 const MAX_BAR_WIDTH: usize = 50; let total = values.len();
729 let below_cutoff = values
730 .iter()
731 .filter(|&&x| below_nnz_cutoff(x, cutoff))
732 .count();
733 let pct_removed = pct(below_cutoff, total);
734
735 let min = values.iter().copied().fold(f32::INFINITY, f32::min);
737 let max = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
738 let sum: f64 = values.iter().map(|&v| v as f64).sum();
740 let mean = if total > 0 { sum / total as f64 } else { 0.0 };
741
742 let mut sorted = values.to_vec();
743 sorted.sort_unstable_by(f32::total_cmp);
744 let median = median_of_sorted(&sorted);
745
746 writeln!(out, "{} {} distribution:", label, metric)?;
747 writeln!(out, " Total: {}", total)?;
748 writeln!(
749 out,
750 " Min: {}, Max: {}, Mean: {:.2}, Median: {:.2}",
751 fmt_stat(min, 2),
752 fmt_stat(max, 2),
753 mean,
754 median
755 )?;
756 if cutoff > 0 {
757 writeln!(
758 out,
759 " Cutoff: {} (removes {} / {} = {:.2}%)",
760 cutoff, below_cutoff, total, pct_removed
761 )?;
762 }
763 if let Some(s) = suggested {
764 let below_s = values.iter().filter(|&&x| below_nnz_cutoff(x, s)).count();
765 let pct_s = pct(below_s, total);
766 writeln!(
767 out,
768 " Suggested cutoff (histogram trough of {}): {} (would remove {} / {} = {:.2}%)",
769 metric, s, below_s, total, pct_s
770 )?;
771 }
772
773 let hist = create_log_histogram(values, cutoff);
775
776 let max_log_count = hist
778 .iter()
779 .map(|b| ((b.count as f64) + 1.0).log10())
780 .fold(0.0_f64, f64::max)
781 .max(1e-9);
782
783 writeln!(
784 out,
785 " Histogram (x: {m} in log-spaced bins, each its range; bar: log10(count+1)):",
786 m = metric
787 )?;
788 for b in hist {
789 let marker = if b.is_cutoff { " <-- CUTOFF" } else { "" };
790 let log_count = ((b.count as f64) + 1.0).log10();
791 let bar_width = ((log_count / max_log_count) * MAX_BAR_WIDTH as f64).round() as usize;
792 let bar_width = if b.count > 0 { bar_width.max(1) } else { 0 };
793 let bar = "█".repeat(bar_width);
794 let range = if b.val_min == b.val_max {
795 fmt_sig(b.val_min)
796 } else {
797 format!("{}-{}", fmt_sig(b.val_min), fmt_sig(b.val_max))
798 };
799 writeln!(out, " {:>17}: {:>6} {}{}", range, b.count, bar, marker)?;
800 }
801 Ok(())
802}
803
804pub fn squeeze_in_place(
809 backend_file: &str,
810 row: usize,
811 column: usize,
812 block_size: Option<usize>,
813) -> anyhow::Result<()> {
814 log::info!("Squeeze the backend data {backend_file}");
815 let (backend, file) = crate::hdf5_io::resolve_backend_file(backend_file, None)?;
816 let data = crate::sparse_io::open_sparse_matrix(&file, &backend)?;
817 squeeze_by_nnz(
818 data.as_ref(),
819 SqueezeCutoffs { row, column },
820 block_size,
821 true,
822 )?;
823 let data = crate::sparse_io::open_sparse_matrix(&file, &backend)?;
824 log::info!(
825 "after squeeze -- data: {} rows x {} columns",
826 data.num_rows().unwrap_or_default(),
827 data.num_columns().unwrap_or_default()
828 );
829 Ok(())
830}
831
832#[cfg(test)]
833mod tests;