1#![allow(dead_code)]
2
3use crate::alg::batch_mixing::*;
4use crate::sparse_data_visitors::*;
5use crate::sparse_io_stack::SparseIoStack;
6use crate::sparse_io_vector::SparseIoVec;
7use legume_numeric::matrix::dmatrix_util::*;
8use legume_numeric::matrix::rand_util::mix_seed;
9use std::sync::{Arc, Mutex};
10
11use legume_numeric::matrix::traits::*;
12use log::{info, warn};
13use nalgebra::DVector;
14
15struct ProjVisitorIn<'a> {
21 basis_kd: &'a nalgebra::DMatrix<f32>,
22}
23
24pub struct RandColProjOut {
25 pub basis: nalgebra::DMatrix<f32>,
26 pub proj: nalgebra::DMatrix<f32>,
27}
28
29pub struct RandRowProjOut {
30 pub basis: nalgebra::DMatrix<f32>,
31 pub proj: nalgebra::DMatrix<f32>,
32}
33
34pub const DEFAULT_PROJECTION_SEED: u64 = 0x50524F4A_50524F4A; pub trait RandProjOps {
44 fn project_columns(
60 &self,
61 target_dim: usize,
62 block_size: Option<usize>,
63 ) -> anyhow::Result<RandColProjOut> {
64 self.project_columns_with_batch_correction::<usize>(target_dim, block_size, None)
65 }
66
67 fn project_columns_with_batch_correction<T>(
83 &self,
84 target_dim: usize,
85 block_size: Option<usize>,
86 batch_membership: Option<&[T]>,
87 ) -> anyhow::Result<RandColProjOut>
88 where
89 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
90 {
91 self.project_columns_with_batch_correction_seeded(
92 target_dim,
93 block_size,
94 batch_membership,
95 DEFAULT_PROJECTION_SEED,
96 )
97 }
98
99 fn project_columns_with_batch_correction_seeded<T>(
102 &self,
103 target_dim: usize,
104 block_size: Option<usize>,
105 batch_membership: Option<&[T]>,
106 seed: u64,
107 ) -> anyhow::Result<RandColProjOut>
108 where
109 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString;
110
111 fn project_columns_weighted<T>(
118 &self,
119 target_dim: usize,
120 block_size: Option<usize>,
121 batch_membership: Option<&[T]>,
122 row_weights: &[f32],
123 ) -> anyhow::Result<RandColProjOut>
124 where
125 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
126 {
127 self.project_columns_weighted_seeded(
128 target_dim,
129 block_size,
130 batch_membership,
131 row_weights,
132 DEFAULT_PROJECTION_SEED,
133 )
134 }
135
136 fn project_columns_weighted_seeded<T>(
139 &self,
140 target_dim: usize,
141 block_size: Option<usize>,
142 batch_membership: Option<&[T]>,
143 row_weights: &[f32],
144 seed: u64,
145 ) -> anyhow::Result<RandColProjOut>
146 where
147 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString;
148
149 fn partition_columns_to_groups(
158 &mut self,
159 proj_kn: &nalgebra::DMatrix<f32>,
160 num_features: Option<usize>,
161 ncols_per_group: Option<usize>,
162 ) -> anyhow::Result<usize>;
163
164 fn partition_columns_to_mixed_groups<T>(
170 &mut self,
171 proj_kn: &nalgebra::DMatrix<f32>,
172 num_features: Option<usize>,
173 batch_membership: &[T],
174 min_batches: usize,
175 merge_levels: usize,
176 ) -> anyhow::Result<usize>
177 where
178 T: std::hash::Hash + Eq + Clone,
179 {
180 let part = mixed_partition(
181 proj_kn,
182 num_features,
183 batch_membership,
184 min_batches,
185 merge_levels,
186 )?;
187 self.assign_group_labels(&part.labels);
188 Ok(part.num_groups)
189 }
190
191 fn assign_group_labels(&mut self, labels: &[usize]);
193
194 }
197
198fn project_columns_visitor(
200 job: (usize, usize),
201 data_vec: &SparseIoVec,
202 shared_in: &ProjVisitorIn<'_>,
203 arc_proj_kn: Arc<Mutex<&mut nalgebra::DMatrix<f32>>>,
204) -> anyhow::Result<()> {
205 let (lb, ub) = job;
206 let basis_kd = shared_in.basis_kd;
207 let target_dim = basis_kd.nrows();
208 let n_block = ub - lb;
209
210 let mut xx_dm = data_vec.read_columns_csc(lb..ub)?;
211 for x in xx_dm.values_mut() {
212 *x = x.ln_1p();
213 }
214 xx_dm.normalize_columns_inplace();
215
216 let mut chunk = nalgebra::DMatrix::<f32>::zeros(target_dim, n_block);
220 for (j, col) in xx_dm.col_iter().enumerate() {
221 for (&i, &v) in col.row_indices().iter().zip(col.values()) {
222 chunk.column_mut(j).axpy(v, &basis_kd.column(i), 1.0);
223 }
224 }
225
226 let mut proj_kn = arc_proj_kn.lock().expect("proj_kn lock");
227 proj_kn.columns_range_mut(lb..ub).copy_from(&chunk);
228 Ok(())
229}
230
231impl RandProjOps for SparseIoStack {
232 fn project_columns_with_batch_correction_seeded<T>(
233 &self,
234 target_dim: usize,
235 block_size: Option<usize>,
236 batch_membership: Option<&[T]>,
237 seed: u64,
238 ) -> anyhow::Result<RandColProjOut>
239 where
240 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
241 {
242 let ncols = self.num_columns()?;
243 let batch_membership = batch_membership.and_then(|x| x.get(0..ncols));
244
245 let proj_vec = self
246 .stack
247 .iter()
248 .enumerate()
249 .map(|(m, data_vec)| -> anyhow::Result<_> {
250 data_vec.project_columns_with_batch_correction_seeded(
253 target_dim,
254 block_size,
255 batch_membership,
256 mix_seed(seed, m as u64),
257 )
258 })
259 .collect::<anyhow::Result<Vec<_>>>()?;
260
261 let basis = concatenate_vertical(
262 &proj_vec
263 .iter()
264 .map(|x| {
265 x.basis.to_owned()
267 })
268 .collect::<Vec<_>>(),
269 )?;
270
271 let proj = concatenate_vertical(
272 &proj_vec
273 .iter()
274 .map(|x| {
275 x.proj.to_owned()
277 })
278 .collect::<Vec<_>>(),
279 )?;
280
281 Ok(RandColProjOut { basis, proj })
282 }
283
284 fn project_columns_weighted_seeded<T>(
285 &self,
286 target_dim: usize,
287 block_size: Option<usize>,
288 batch_membership: Option<&[T]>,
289 row_weights: &[f32],
290 seed: u64,
291 ) -> anyhow::Result<RandColProjOut>
292 where
293 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
294 {
295 let ncols = self.num_columns()?;
296 let batch_membership = batch_membership.and_then(|x| x.get(0..ncols));
297
298 let total_rows: usize = self.stack.iter().map(|v| v.num_rows()).sum();
303 if row_weights.len() != total_rows {
304 return Err(anyhow::anyhow!(
305 "row_weights length {} does not match total feature count {} \
306 (sum of num_rows() across the stack)",
307 row_weights.len(),
308 total_rows
309 ));
310 }
311
312 let mut offset = 0usize;
313 let mut proj_vec = Vec::with_capacity(self.stack.len());
314 for (m, data_vec) in self.stack.iter().enumerate() {
315 let nrows = data_vec.num_rows();
316 let slice = &row_weights[offset..offset + nrows];
317 offset += nrows;
318 proj_vec.push(data_vec.project_columns_weighted_seeded(
319 target_dim,
320 block_size,
321 batch_membership,
322 slice,
323 mix_seed(seed, m as u64),
324 )?);
325 }
326
327 let basis = concatenate_vertical(
328 &proj_vec
329 .iter()
330 .map(|x| x.basis.to_owned())
331 .collect::<Vec<_>>(),
332 )?;
333 let proj = concatenate_vertical(
334 &proj_vec
335 .iter()
336 .map(|x| x.proj.to_owned())
337 .collect::<Vec<_>>(),
338 )?;
339
340 Ok(RandColProjOut { basis, proj })
341 }
342
343 fn partition_columns_to_groups(
344 &mut self,
345 proj_kn: &nalgebra::DMatrix<f32>,
346 num_sorting_features: Option<usize>,
347 ncols_per_group: Option<usize>,
348 ) -> anyhow::Result<usize> {
349 let nn = proj_kn.ncols();
350 if nn != self.num_columns()? {
351 return Err(anyhow::anyhow!("number of columns mismatch"));
352 }
353
354 let target_kk = num_sorting_features.unwrap_or(proj_kn.nrows());
355 let kk = proj_kn.nrows().min(target_kk).min(nn);
356
357 let binary_codes = binary_sort_columns(proj_kn, kk)?;
358 let max_group = *binary_codes
359 .iter()
360 .max()
361 .ok_or(anyhow::anyhow!("unable to determine max element"))?;
362
363 for x in self.stack.iter_mut() {
364 x.assign_groups(&binary_codes, ncols_per_group);
365 }
366
367 Ok(max_group + 1)
368 }
369
370 fn assign_group_labels(&mut self, labels: &[usize]) {
371 for x in self.stack.iter_mut() {
372 x.assign_groups(labels, None);
373 }
374 }
375}
376
377impl RandProjOps for SparseIoVec {
378 fn project_columns_with_batch_correction_seeded<T>(
379 &self,
380 target_dim: usize,
381 block_size: Option<usize>,
382 batch_membership: Option<&[T]>,
383 seed: u64,
384 ) -> anyhow::Result<RandColProjOut>
385 where
386 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
387 {
388 let nrows = self.num_rows();
389 let ncols = self.num_columns();
390
391 let mut proj_kn = nalgebra::DMatrix::<f32>::zeros(target_dim, ncols);
392
393 let basis_dk = nalgebra::DMatrix::<f32>::rnorm_seeded(nrows, target_dim, seed);
394 let basis_kd = basis_dk.transpose();
397
398 let shared_in = ProjVisitorIn {
399 basis_kd: &basis_kd,
400 };
401 info!(
402 "starting random projection loop (project_columns_with_batch_correction): \
403 D={nrows}, N={ncols}, K={target_dim}, block_size={:?}",
404 block_size
405 );
406 self.visit_columns_by_block(
407 &project_columns_visitor,
408 &shared_in,
409 &mut proj_kn,
410 block_size,
411 )?;
412 info!("finished random projection loop");
413
414 if let Some(col_to_batch) = batch_membership {
415 adjust_batch_biases(&mut proj_kn, col_to_batch)?;
416 }
417
418 let (lb, ub) = (-4., 4.);
419 proj_kn.scale_columns_inplace();
420
421 if proj_kn.max() > ub || proj_kn.min() < lb {
422 proj_kn.iter_mut().for_each(|x| {
424 *x = x.clamp(lb, ub);
425 });
426 proj_kn.scale_columns_inplace();
427 }
428 Ok(RandColProjOut {
429 basis: basis_dk,
430 proj: proj_kn,
431 })
432 }
433
434 fn project_columns_weighted_seeded<T>(
438 &self,
439 target_dim: usize,
440 block_size: Option<usize>,
441 batch_membership: Option<&[T]>,
442 row_weights: &[f32],
443 seed: u64,
444 ) -> anyhow::Result<RandColProjOut>
445 where
446 T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
447 {
448 let nrows = self.num_rows();
449 let ncols = self.num_columns();
450
451 assert_eq!(row_weights.len(), nrows, "row_weights length mismatch");
452
453 let mut proj_kn = nalgebra::DMatrix::<f32>::zeros(target_dim, ncols);
454
455 let mut basis_dk = nalgebra::DMatrix::<f32>::rnorm_seeded(nrows, target_dim, seed);
456
457 for (r, &w) in row_weights.iter().enumerate() {
459 if w <= 0.0 {
460 basis_dk.row_mut(r).fill(0.0);
461 } else if (w - 1.0).abs() > 1e-6 {
462 basis_dk.row_mut(r).scale_mut(w);
463 }
464 }
465
466 let basis_kd = basis_dk.transpose();
467 let shared_in = ProjVisitorIn {
468 basis_kd: &basis_kd,
469 };
470 info!(
471 "starting random projection loop (project_columns_weighted): \
472 D={nrows}, N={ncols}, K={target_dim}, block_size={:?}",
473 block_size
474 );
475 self.visit_columns_by_block(
476 &project_columns_visitor,
477 &shared_in,
478 &mut proj_kn,
479 block_size,
480 )?;
481 info!("finished random projection loop");
482
483 if let Some(col_to_batch) = batch_membership {
484 adjust_batch_biases(&mut proj_kn, col_to_batch)?;
485 }
486
487 let (lb, ub) = (-4., 4.);
488 proj_kn.scale_columns_inplace();
489 if proj_kn.max() > ub || proj_kn.min() < lb {
490 proj_kn.iter_mut().for_each(|x| {
491 *x = x.clamp(lb, ub);
492 });
493 proj_kn.scale_columns_inplace();
494 }
495 Ok(RandColProjOut {
496 basis: basis_dk,
497 proj: proj_kn,
498 })
499 }
500
501 fn partition_columns_to_groups(
511 &mut self,
512 proj_kn: &nalgebra::DMatrix<f32>,
513 num_sorting_features: Option<usize>,
514 ncols_per_group: Option<usize>,
515 ) -> anyhow::Result<usize> {
516 let nn = proj_kn.ncols();
517 if nn != self.num_columns() {
518 return Err(anyhow::anyhow!("number of columns mismatch"));
519 }
520
521 let target_kk = num_sorting_features.unwrap_or(proj_kn.nrows());
522 let kk = proj_kn.nrows().min(target_kk).min(nn);
523
524 let binary_codes = binary_sort_columns(proj_kn, kk)?;
525 let max_group = *binary_codes
526 .iter()
527 .max()
528 .ok_or(anyhow::anyhow!("unable to determine max element"))?;
529 self.assign_groups(&binary_codes, ncols_per_group);
530 Ok(max_group + 1)
531 }
532
533 fn assign_group_labels(&mut self, labels: &[usize]) {
534 self.assign_groups(labels, None);
535 }
536}
537
538fn adjust_batch_biases<T>(
543 proj_kn: &mut nalgebra::DMatrix<f32>,
544 col_to_batch: &[T],
545) -> anyhow::Result<()>
546where
547 T: std::hash::Hash + Eq + Clone,
548{
549 let ncols = proj_kn.ncols();
550 if col_to_batch.len() != ncols {
551 warn!(
552 "The batch membership size {} mismatches with the number of columns {} ...",
553 col_to_batch.len(),
554 ncols
555 );
556 return Ok(());
557 }
558 let batch = batch_indices(col_to_batch);
559 let n_batches = batch.iter().max().map_or(0, |&m| m + 1);
560 let bits = default_state_bits(ncols, n_batches, proj_kn.nrows());
561 info!(
562 "adjusting batch biases within cell state ({} batches, {} state bits) ...",
563 n_batches, bits
564 );
565 *proj_kn = centre_batches_within_state(proj_kn, &batch, bits)?;
566 Ok(())
567}
568
569pub struct MixedPartition {
572 pub codes: Vec<usize>,
574 pub labels: Vec<usize>,
576 pub num_groups: usize,
577 pub bits: usize,
579 pub batch: Vec<usize>,
581 pub min_batches: usize,
583}
584
585pub fn mixed_partition<T>(
589 proj_kn: &nalgebra::DMatrix<f32>,
590 num_features: Option<usize>,
591 batch_membership: &[T],
592 min_batches: usize,
593 merge_levels: usize,
594) -> anyhow::Result<MixedPartition>
595where
596 T: std::hash::Hash + Eq + Clone,
597{
598 let nn = proj_kn.ncols();
599 anyhow::ensure!(
600 batch_membership.len() == nn,
601 "batch membership size {} mismatches the number of columns {}",
602 batch_membership.len(),
603 nn
604 );
605 let batch = batch_indices(batch_membership);
606 let n_batches = batch.iter().max().map_or(0, |&m| m + 1);
607 let min_batches = min_batches.min(n_batches);
608 let bits = proj_kn
609 .nrows()
610 .min(num_features.unwrap_or(proj_kn.nrows()))
611 .min(nn);
612 let codes = binary_sort_columns(proj_kn, bits)?;
613 let labels = merge_poorly_mixed_bins(&codes, &batch, bits, merge_levels, min_batches);
614 let num_groups = labels
615 .iter()
616 .collect::<std::collections::HashSet<_>>()
617 .len();
618 info!(
619 "partitioned columns into {} groups ({} bits, bins with fewer than {} batches merged up to {} levels)",
620 num_groups, bits, min_batches, merge_levels
621 );
622 Ok(MixedPartition {
623 codes,
624 labels,
625 num_groups,
626 bits,
627 batch,
628 min_batches,
629 })
630}
631
632pub fn binary_sort_columns(
638 proj_kn: &nalgebra::DMatrix<f32>,
639 kk: usize,
640) -> anyhow::Result<Vec<usize>> {
641 let nn = proj_kn.ncols();
650 let (_, _, mut q_nk) = proj_kn.rsvd(kk)?;
651 q_nk.scale_columns_inplace();
652
653 let mut binary_codes = DVector::<usize>::zeros(nn);
654 for k in 0..kk {
655 let binary_shift = |x: f32| -> usize {
656 if x > 0.0 {
657 1 << k
658 } else {
659 0
660 }
661 };
662 binary_codes += q_nk.column(k).map(binary_shift);
663 }
664
665 Ok(binary_codes.data.as_vec().clone())
666}
667
668pub fn max_binary_code(kk: usize) -> usize {
670 (1 << kk) - 1
671}
672
673#[cfg(test)]
674mod repro_tests {
675 use super::*;
676 use crate::sparse_io::{create_sparse_from_triplets, SparseIoBackend};
677 use std::sync::Arc;
678
679 fn tiny_data_vec(dir: &tempfile::TempDir) -> anyhow::Result<SparseIoVec> {
681 let path = dir.path().join("proj.zarr");
682 let path_str = path.to_str().unwrap();
683 let (d, n) = (8usize, 12usize);
684 let mut triplets: Vec<(u64, u64, f32)> = Vec::new();
686 for i in 0..d {
687 for j in 0..n {
688 if (i * 7 + j * 3) % 5 < 3 {
689 triplets.push((i as u64, j as u64, 1.0 + ((i + j) % 4) as f32));
690 }
691 }
692 }
693 let nnz = triplets.len();
694 let mut backend = create_sparse_from_triplets(
695 &triplets,
696 (d, n, nnz),
697 Some(path_str),
698 Some(&SparseIoBackend::Zarr),
699 )?;
700 let gene_names: Vec<Box<str>> = (0..d).map(|i| format!("g{i}").into()).collect();
701 let cell_names: Vec<Box<str>> = (0..n).map(|j| format!("c{j}").into()).collect();
702 backend.register_row_names_vec(&gene_names);
703 backend.register_column_names_vec(&cell_names);
704 let mut data_vec = SparseIoVec::new();
705 data_vec.push(Arc::from(backend), None)?;
706 Ok(data_vec)
707 }
708
709 #[test]
712 fn seeded_projection_is_reproducible() -> anyhow::Result<()> {
713 let dir = tempfile::tempdir()?;
714 let data_vec = tiny_data_vec(&dir)?;
715 let a =
716 data_vec.project_columns_with_batch_correction_seeded::<usize>(4, None, None, 123)?;
717 let b =
718 data_vec.project_columns_with_batch_correction_seeded::<usize>(4, None, None, 123)?;
719 assert_eq!(a.basis, b.basis, "same seed → identical basis");
720 assert_eq!(a.proj, b.proj, "same seed → identical projection");
721 Ok(())
722 }
723
724 #[test]
727 fn default_projection_is_reproducible() -> anyhow::Result<()> {
728 let dir = tempfile::tempdir()?;
729 let data_vec = tiny_data_vec(&dir)?;
730 let a = data_vec.project_columns(4, None)?;
731 let b = data_vec.project_columns(4, None)?;
732 assert_eq!(a.basis, b.basis);
733 assert_eq!(a.proj, b.proj);
734 Ok(())
735 }
736
737 #[test]
740 fn distinct_seeds_diverge() -> anyhow::Result<()> {
741 let dir = tempfile::tempdir()?;
742 let data_vec = tiny_data_vec(&dir)?;
743 let a = data_vec.project_columns_with_batch_correction_seeded::<usize>(4, None, None, 1)?;
744 let b = data_vec.project_columns_with_batch_correction_seeded::<usize>(4, None, None, 2)?;
745 assert_ne!(a.basis, b.basis, "distinct seeds → distinct basis");
746 Ok(())
747 }
748}