1use rayon::prelude::*;
15
16use crate::limma::lm_fit::lm_fit;
17use crate::prelude::*;
18
19fn sum_contrast_columns(labels: &[usize]) -> Result<(Vec<f64>, usize), EdgeErrors> {
39 let mut levels = labels.to_vec();
40 levels.sort_unstable();
41 levels.dedup();
42 if levels.len() < 2 {
43 return Err(EdgeErrors::InvalidArgument(
44 "a batch factor needs at least two levels".to_string(),
45 ));
46 }
47
48 let n_cols = levels.len() - 1;
49 let mut out = vec![0.0; labels.len() * n_cols];
50 for (row, label) in out.chunks_exact_mut(n_cols).zip(labels) {
51 let lvl = levels.binary_search(label).expect("label is a level");
53 if lvl == n_cols {
54 row.fill(-1.0);
55 } else {
56 row[lvl] = 1.0;
57 }
58 }
59 Ok((out, n_cols))
60}
61
62fn cbind(left: &[f64], n_left: usize, right: &[f64], n_right: usize, n_rows: usize) -> Vec<f64> {
76 let mut out = Vec::with_capacity(n_rows * (n_left + n_right));
77 for i in 0..n_rows {
78 out.extend_from_slice(&left[i * n_left..(i + 1) * n_left]);
79 out.extend_from_slice(&right[i * n_right..(i + 1) * n_right]);
80 }
81 out
82}
83
84fn check_len(name: &'static str, len: usize, n_samples: usize) -> Result<(), EdgeErrors> {
96 if len != n_samples {
97 return Err(EdgeErrors::LengthMismatch {
98 name,
99 expected: n_samples,
100 got: len,
101 });
102 }
103 Ok(())
104}
105
106#[allow(clippy::too_many_arguments)]
143pub fn remove_batch_effect<T: EdgeFloat>(
144 x: &[T],
145 n_genes: usize,
146 n_samples: usize,
147 batch: Option<&[usize]>,
148 batch2: Option<&[usize]>,
149 covariates: Option<(&[f64], usize)>,
150 design: Option<(&[f64], usize)>,
151 weights: Option<&Recycled<f64>>,
152) -> Result<Vec<f64>, EdgeErrors> {
153 if x.len() != n_genes * n_samples {
154 return Err(EdgeErrors::LengthMismatch {
155 name: "x",
156 expected: n_genes * n_samples,
157 got: x.len(),
158 });
159 }
160 let y: Vec<f64> = x.iter().map(|v| v.to_f64().unwrap_or(f64::NAN)).collect();
161
162 let mut x_batch: Vec<f64> = vec![0.0; 0];
163 let mut n_batch = 0;
164 for (name, labels) in [("batch", batch), ("batch2", batch2)] {
165 if let Some(labels) = labels {
166 check_len(name, labels.len(), n_samples)?;
167 let (cols, n_cols) = sum_contrast_columns(labels)?;
168 x_batch = cbind(&x_batch, n_batch, &cols, n_cols, n_samples);
169 n_batch += n_cols;
170 }
171 }
172 if let Some((cov, n_cov)) = covariates {
173 check_len("covariates", cov.len(), n_samples * n_cov)?;
174 let mut centred = cov.to_vec();
175 for j in 0..n_cov {
176 let mean = (0..n_samples).map(|i| cov[i * n_cov + j]).sum::<f64>() / n_samples as f64;
177 (0..n_samples).for_each(|i| centred[i * n_cov + j] -= mean);
178 }
179 x_batch = cbind(&x_batch, n_batch, ¢red, n_cov, n_samples);
180 n_batch += n_cov;
181 }
182 if n_batch == 0 {
183 return Ok(y);
184 }
185
186 let intercept = vec![1.0; n_samples];
187 let (design, n_coef) = design.unwrap_or((&intercept, 1));
188 check_len("design", design.len(), n_samples * n_coef)?;
189
190 let n_full = n_coef + n_batch;
191 let full = cbind(design, n_coef, &x_batch, n_batch, n_samples);
192 let fit = lm_fit(&y, n_genes, n_samples, &full, n_full, weights, None, None)?;
193
194 let mut out = y;
195 out.par_chunks_mut(n_samples)
196 .zip(fit.coefficients.par_chunks(n_full))
197 .for_each(|(row, coef)| {
198 let beta = &coef[n_coef..];
199 for (s, v) in row.iter_mut().enumerate() {
200 let x_s = &x_batch[s * n_batch..(s + 1) * n_batch];
201 let effect: f64 = beta
203 .iter()
204 .zip(x_s)
205 .map(|(b, xb)| if b.is_nan() { 0.0 } else { b * xb })
206 .sum();
207 *v -= effect;
208 }
209 });
210 Ok(out)
211}
212
213#[cfg(test)]
218mod tests {
219 #![allow(clippy::excessive_precision)]
223
224 use super::*;
225 use approx::assert_relative_eq;
226
227 const N_GENES: usize = 8;
228 const N_SAMPLES: usize = 6;
229
230 const TOL: f64 = 1e-13;
233
234 fn fixture_y() -> Vec<f64> {
237 (0..N_GENES * N_SAMPLES)
238 .map(|i| ((i * 7) % 23) as f64 / 64.0)
239 .collect()
240 }
241
242 const BATCH: [usize; 6] = [0, 0, 1, 1, 2, 2];
243 const BATCH2: [usize; 6] = [0, 1, 0, 1, 1, 0];
244
245 const DESIGN: [f64; 12] = [1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0];
247
248 const COV: [f64; 6] = [0.125, 0.375, 0.25, 0.625, 0.5, 0.875];
250
251 fn assert_matches(got: &[f64], expected: &[f64]) {
252 assert_eq!(got.len(), expected.len());
253 for (g, e) in got.iter().zip(expected) {
254 assert_relative_eq!(*g, *e, epsilon = TOL);
255 }
256 }
257
258 #[test]
259 fn test_remove_batch_effect_batch_only() {
260 let expected = [
262 0.098958333333333343,
263 0.208333333333333343,
264 0.098958333333333287,
265 0.208333333333333287,
266 0.098958333333333370,
267 0.208333333333333370,
268 0.276041666666666630,
269 0.026041666666666657,
270 0.096354166666666657,
271 0.205729166666666657,
272 0.096354166666666685,
273 0.205729166666666685,
274 0.153645833333333287,
275 0.263020833333333259,
276 0.153645833333333370,
277 0.263020833333333370,
278 0.333333333333333370,
279 0.083333333333333343,
280 0.091145833333333287,
281 0.200520833333333287,
282 0.091145833333333370,
283 0.200520833333333370,
284 0.270833333333333315,
285 0.020833333333333336,
286 0.148437499999999972,
287 0.257812500000000000,
288 0.328125000000000000,
289 0.078125000000000000,
290 0.148437500000000028,
291 0.257812500000000000,
292 0.085937500000000000,
293 0.195312500000000000,
294 0.265625000000000000,
295 0.015625000000000007,
296 0.085937500000000000,
297 0.195312500000000000,
298 0.322916666666666630,
299 0.072916666666666630,
300 0.143229166666666685,
301 0.252604166666666685,
302 0.143229166666666685,
303 0.252604166666666685,
304 0.260416666666666685,
305 0.010416666666666661,
306 0.080729166666666657,
307 0.190104166666666657,
308 0.080729166666666685,
309 0.190104166666666685,
310 ];
311 let got = remove_batch_effect(
312 &fixture_y(),
313 N_GENES,
314 N_SAMPLES,
315 Some(&BATCH),
316 None,
317 None,
318 None,
319 None,
320 )
321 .unwrap();
322 assert_matches(&got, &expected);
323 }
324
325 #[test]
326 fn test_remove_batch_effect_with_design() {
327 let expected = [
329 0.098958333333333329,
330 0.208333333333333315,
331 0.153645833333333315,
332 0.263020833333333315,
333 0.044270833333333370,
334 0.153645833333333370,
335 0.276041666666666685,
336 0.026041666666666689,
337 -0.028645833333333481,
338 0.080729166666666519,
339 0.221354166666666796,
340 0.330729166666666796,
341 0.153645833333333315,
342 0.263020833333333315,
343 0.208333333333333343,
344 0.317708333333333370,
345 0.278645833333333370,
346 0.028645833333333343,
347 0.091145833333333315,
348 0.200520833333333315,
349 0.145833333333333343,
350 0.255208333333333370,
351 0.216145833333333343,
352 -0.033854166666666657,
353 0.148437499999999944,
354 0.257812499999999944,
355 0.382812500000000000,
356 0.132812500000000000,
357 0.093750000000000042,
358 0.203125000000000056,
359 0.085937499999999958,
360 0.195312499999999944,
361 0.320312500000000056,
362 0.070312500000000056,
363 0.031249999999999986,
364 0.140625000000000000,
365 0.322916666666666685,
366 0.072916666666666671,
367 0.018229166666666519,
368 0.127604166666666519,
369 0.268229166666666852,
370 0.377604166666666852,
371 0.260416666666666685,
372 0.010416666666666678,
373 -0.044270833333333426,
374 0.065104166666666574,
375 0.205729166666666741,
376 0.315104166666666741,
377 ];
378 let got = remove_batch_effect(
379 &fixture_y(),
380 N_GENES,
381 N_SAMPLES,
382 Some(&BATCH),
383 None,
384 None,
385 Some((&DESIGN, 2)),
386 None,
387 )
388 .unwrap();
389 assert_matches(&got, &expected);
390 }
391
392 #[test]
393 fn test_remove_batch_effect_two_batches_and_design() {
394 let expected = [
396 0.0989583333333333148,
397 0.2083333333333333426,
398 0.1536458333333333148,
399 0.2630208333333333148,
400 0.0442708333333333634,
401 0.1536458333333333426,
402 0.2760416666666666852,
403 0.0260416666666666644,
404 -0.0286458333333334814,
405 0.0807291666666664631,
406 0.2213541666666667962,
407 0.3307291666666668517,
408 0.2434895833333333148,
409 0.1731770833333333148,
410 0.2083333333333332871,
411 0.1380208333333332871,
412 0.2786458333333333703,
413 0.2083333333333333981,
414 0.1809895833333333148,
415 0.1106770833333333148,
416 0.1458333333333332871,
417 0.0755208333333332593,
418 0.2161458333333333981,
419 0.1458333333333334259,
420 0.0585937499999999514,
421 0.3476562499999999445,
422 0.3828125000000000555,
423 0.3125000000000000555,
424 0.0937500000000000000,
425 0.0234375000000000000,
426 -0.0039062500000000486,
427 0.2851562500000000000,
428 0.3203125000000001110,
429 0.2500000000000001110,
430 0.0312499999999999306,
431 -0.0390625000000001110,
432 0.3229166666666666852,
433 0.0729166666666666713,
434 0.0182291666666665186,
435 0.1276041666666665186,
436 0.2682291666666668517,
437 0.3776041666666668517,
438 0.2604166666666666852,
439 0.0104166666666666748,
440 -0.0442708333333334259,
441 0.0651041666666665741,
442 0.2057291666666667407,
443 0.3151041666666667407,
444 ];
445 let got = remove_batch_effect(
446 &fixture_y(),
447 N_GENES,
448 N_SAMPLES,
449 Some(&BATCH),
450 Some(&BATCH2),
451 None,
452 Some((&DESIGN, 2)),
453 None,
454 )
455 .unwrap();
456 assert_matches(&got, &expected);
457 }
458
459 #[test]
460 fn test_remove_batch_effect_covariates_and_design() {
461 let expected = [
463 0.135416666666666657,
464 0.171875000000000000,
465 0.171874999999999972,
466 0.171874999999999972,
467 0.135416666666666657,
468 0.135416666666666685,
469 0.312500000000000056,
470 -0.010416666666666650,
471 -0.010416666666666824,
472 -0.010416666666666852,
473 0.312500000000000111,
474 0.312500000000000111,
475 0.130208333333333315,
476 0.286458333333333315,
477 0.196614583333333343,
478 0.376302083333333370,
479 0.220052083333333343,
480 0.040364583333333343,
481 0.067708333333333315,
482 0.223958333333333343,
483 0.134114583333333315,
484 0.313802083333333370,
485 0.157552083333333315,
486 -0.022135416666666647,
487 0.124999999999999944,
488 0.281249999999999944,
489 0.371093750000000000,
490 0.191406250000000028,
491 0.035156250000000028,
492 0.214843750000000028,
493 0.062499999999999958,
494 0.218749999999999944,
495 0.308593750000000056,
496 0.128906250000000056,
497 -0.027343750000000000,
498 0.152343750000000000,
499 0.359375000000000000,
500 0.036458333333333329,
501 0.036458333333333176,
502 0.036458333333333148,
503 0.359375000000000167,
504 0.359375000000000167,
505 0.296875000000000000,
506 -0.026041666666666657,
507 -0.026041666666666768,
508 -0.026041666666666796,
509 0.296875000000000111,
510 0.296875000000000111,
511 ];
512 let got = remove_batch_effect(
513 &fixture_y(),
514 N_GENES,
515 N_SAMPLES,
516 Some(&BATCH),
517 None,
518 Some((&COV, 1)),
519 Some((&DESIGN, 2)),
520 None,
521 )
522 .unwrap();
523 assert_matches(&got, &expected);
524 }
525
526 #[test]
527 fn test_remove_batch_effect_aliased_batch_is_zeroed() {
528 let batch = [0, 0, 0, 1, 1, 1];
531 let design = [1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
532 let y = fixture_y();
533 let got = remove_batch_effect(
534 &y,
535 N_GENES,
536 N_SAMPLES,
537 Some(&batch),
538 None,
539 None,
540 Some((&design, 2)),
541 None,
542 )
543 .unwrap();
544 assert_matches(&got, &y);
545 }
546
547 #[test]
548 fn test_remove_batch_effect_non_contiguous_labels() {
549 let expected = [
551 0.091145833333333329,
552 0.200520833333333315,
553 0.127604166666666685,
554 0.236979166666666685,
555 0.169270833333333315,
556 0.096354166666666671,
557 0.328125000000000000,
558 0.078125000000000000,
559 0.125000000000000000,
560 0.234375000000000000,
561 0.046875000000000000,
562 0.093750000000000000,
563 0.145833333333333315,
564 0.255208333333333315,
565 0.182291666666666685,
566 0.291666666666666685,
567 0.223958333333333315,
568 0.151041666666666685,
569 0.083333333333333329,
570 0.192708333333333315,
571 0.119791666666666671,
572 0.229166666666666685,
573 0.161458333333333315,
574 0.088541666666666671,
575 0.140624999999999972,
576 0.249999999999999972,
577 0.296875000000000000,
578 0.046875000000000028,
579 0.218749999999999972,
580 0.265625000000000000,
581 0.078124999999999986,
582 0.187500000000000000,
583 0.234375000000000000,
584 -0.015624999999999990,
585 0.156250000000000000,
586 0.203125000000000000,
587 0.375000000000000000,
588 0.124999999999999972,
589 0.171875000000000028,
590 0.281250000000000000,
591 0.093749999999999972,
592 0.140625000000000028,
593 0.312500000000000000,
594 0.062500000000000000,
595 0.109375000000000000,
596 0.218750000000000000,
597 0.031250000000000000,
598 0.078125000000000000,
599 ];
600 let batch = [7, 7, 3, 3, 7, 3];
601 let got = remove_batch_effect(
602 &fixture_y(),
603 N_GENES,
604 N_SAMPLES,
605 Some(&batch),
606 None,
607 None,
608 None,
609 None,
610 )
611 .unwrap();
612 assert_matches(&got, &expected);
613 }
614
615 #[test]
616 fn test_remove_batch_effect_nothing_to_remove_returns_input() {
617 let y = fixture_y();
618 let got =
619 remove_batch_effect(&y, N_GENES, N_SAMPLES, None, None, None, None, None).unwrap();
620 assert_matches(&got, &y);
621 }
622
623 #[test]
624 fn test_remove_batch_effect_f32_input() {
625 let y32: Vec<f32> = fixture_y().iter().map(|v| *v as f32).collect();
626 let got32 = remove_batch_effect(
627 &y32,
628 N_GENES,
629 N_SAMPLES,
630 Some(&BATCH),
631 None,
632 None,
633 None,
634 None,
635 )
636 .unwrap();
637 let got64 = remove_batch_effect(
638 &fixture_y(),
639 N_GENES,
640 N_SAMPLES,
641 Some(&BATCH),
642 None,
643 None,
644 None,
645 None,
646 )
647 .unwrap();
648 assert_matches(&got32, &got64);
650 }
651
652 #[test]
653 fn test_remove_batch_effect_single_level_errors() {
654 let batch = [1; 6];
655 let res = remove_batch_effect(
656 &fixture_y(),
657 N_GENES,
658 N_SAMPLES,
659 Some(&batch),
660 None,
661 None,
662 None,
663 None,
664 );
665 assert!(matches!(res, Err(EdgeErrors::InvalidArgument(_))));
666 }
667
668 #[test]
669 fn test_remove_batch_effect_length_mismatch_errors() {
670 let batch = [0, 1, 0];
671 let res = remove_batch_effect(
672 &fixture_y(),
673 N_GENES,
674 N_SAMPLES,
675 Some(&batch),
676 None,
677 None,
678 None,
679 None,
680 );
681 assert!(matches!(
682 res,
683 Err(EdgeErrors::LengthMismatch { name: "batch", .. })
684 ));
685 }
686
687 #[test]
688 fn test_remove_batch_effect_nan_input() {
689 let expected = [
691 0.039062500000000000,
692 0.148437500000000000,
693 f64::NAN,
694 f64::NAN,
695 0.039062500000000000,
696 0.148437500000000000,
697 0.276041666666666630,
698 0.026041666666666657,
699 0.096354166666666657,
700 0.205729166666666657,
701 0.096354166666666685,
702 0.205729166666666685,
703 ];
704 let mut y = fixture_y();
705 y[2] = f64::NAN;
706 y[3] = f64::NAN;
707 let got = remove_batch_effect(&y, N_GENES, N_SAMPLES, Some(&BATCH), None, None, None, None)
708 .unwrap();
709 for (g, e) in got[..12].iter().zip(&expected) {
710 if e.is_nan() {
711 assert!(g.is_nan());
712 } else {
713 assert_relative_eq!(*g, *e, epsilon = TOL);
714 }
715 }
716 }
717
718 #[test]
719 fn test_remove_batch_effect_two_covariates() {
720 let expected = [
723 0.135416666666666657,
724 0.171875000000000000,
725 0.171874999999999972,
726 0.171874999999999972,
727 0.135416666666666657,
728 0.135416666666666685,
729 0.312500000000000056,
730 -0.010416666666666685,
731 -0.010416666666666824,
732 -0.010416666666666907,
733 0.312500000000000111,
734 0.312500000000000111,
735 0.220052083333333287,
736 0.196614583333333343,
737 0.196614583333333315,
738 0.196614583333333370,
739 0.220052083333333426,
740 0.220052083333333287,
741 0.157552083333333287,
742 0.134114583333333315,
743 0.134114583333333315,
744 0.134114583333333315,
745 0.157552083333333398,
746 0.157552083333333343,
747 0.035156249999999944,
748 0.371093749999999944,
749 0.371093750000000056,
750 0.371093750000000056,
751 0.035156249999999972,
752 0.035156250000000056,
753 -0.027343750000000056,
754 0.308593750000000000,
755 0.308593750000000111,
756 0.308593750000000111,
757 -0.027343750000000083,
758 -0.027343750000000056,
759 0.359375000000000000,
760 0.036458333333333322,
761 0.036458333333333176,
762 0.036458333333333148,
763 0.359375000000000167,
764 0.359375000000000167,
765 0.296875000000000056,
766 -0.026041666666666678,
767 -0.026041666666666768,
768 -0.026041666666666852,
769 0.296875000000000111,
770 0.296875000000000111,
771 ];
772 let cov = [
773 0.125, 0.5, 0.375, 0.125, 0.25, 0.375, 0.625, 0.0, 0.5, 0.25, 0.875, 0.625,
774 ];
775 let got = remove_batch_effect(
776 &fixture_y(),
777 N_GENES,
778 N_SAMPLES,
779 Some(&BATCH),
780 None,
781 Some((&cov, 2)),
782 Some((&DESIGN, 2)),
783 None,
784 )
785 .unwrap();
786 assert_matches(&got, &expected);
787 }
788
789 #[test]
790 fn test_remove_batch_effect_covariates_and_design_length_errors() {
791 let res = remove_batch_effect(
792 &fixture_y(),
793 N_GENES,
794 N_SAMPLES,
795 None,
796 None,
797 Some((&COV, 2)),
798 None,
799 None,
800 );
801 assert!(matches!(
802 res,
803 Err(EdgeErrors::LengthMismatch {
804 name: "covariates",
805 ..
806 })
807 ));
808
809 let res = remove_batch_effect(
811 &fixture_y(),
812 N_GENES,
813 N_SAMPLES,
814 Some(&BATCH),
815 None,
816 None,
817 Some((&DESIGN, 3)),
818 None,
819 );
820 assert!(matches!(
821 res,
822 Err(EdgeErrors::LengthMismatch { name: "design", .. })
823 ));
824 }
825
826 #[test]
827 fn test_remove_batch_effect_weights() {
828 let expected = [
831 0.093894675925925930,
832 0.203269675925925930,
833 0.157696759259259189,
834 0.267071759259259189,
835 0.045283564814814867,
836 0.154658564814814881,
837 0.279839409722222265,
838 0.029839409722222252,
839 -0.018012152777777873,
840 0.091362847222222127,
841 0.206922743055555636,
842 0.316297743055555636,
843 0.167601495726495797,
844 0.276976495726495797,
845 0.201255341880341887,
846 0.310630341880341887,
847 0.271768162393162316,
848 0.021768162393162316,
849 0.092708333333333323,
850 0.202083333333333337,
851 0.158333333333333326,
852 0.267708333333333326,
853 0.202083333333333337,
854 -0.047916666666666649,
855 0.161401098901098883,
856 0.270776098901098883,
857 0.377918956043956089,
858 0.127918956043956089,
859 0.085679945054945028,
860 0.195054945054945028,
861 0.086921296296296302,
862 0.196296296296296302,
863 0.307407407407407407,
864 0.057407407407407407,
865 0.043171296296296291,
866 0.152546296296296291,
867 0.326388888888888840,
868 0.076388888888888867,
869 0.003472222222222154,
870 0.112847222222222154,
871 0.279513888888888951,
872 0.388888888888888951,
873 0.255353009259259300,
874 0.005353009259259283,
875 -0.040219907407407413,
876 0.069155092592592587,
877 0.206741898148148140,
878 0.316116898148148140,
879 ];
880 let w = (0..N_GENES * N_SAMPLES)
881 .map(|i| 1.0 + ((i * 5) % 7) as f64 / 4.0)
882 .collect();
883 let w = Recycled::full(w, N_GENES, N_SAMPLES).unwrap();
884 let got = remove_batch_effect(
885 &fixture_y(),
886 N_GENES,
887 N_SAMPLES,
888 Some(&BATCH),
889 None,
890 None,
891 Some((&DESIGN, 2)),
892 Some(&w),
893 )
894 .unwrap();
895 assert_matches(&got, &expected);
896 }
897}