Skip to main content

edge_rs/limma/
remove_batch_effect.rs

1//! limma's `removeBatchEffect`: regress batch out of a log-expression matrix.
2//!
3//! Batch factors get sum-to-zero contrasts, so corrected values stay centred on
4//! the grand mean. The batch columns are fitted jointly with the design of
5//! interest through [`lm_fit`] and only the batch part is subtracted, so the
6//! design protects the biology.
7//!
8//! For plotting and unsupervised work. For testing, put batch into the design.
9//!
10//! ### References
11//!
12//! Smyth, Statistical Applications in Genetics and Molecular Biology 3(1), 2004
13
14use rayon::prelude::*;
15
16use crate::limma::lm_fit::lm_fit;
17use crate::prelude::*;
18
19/////////////
20// Helpers //
21/////////////
22
23/// Sum-to-zero encoding of one batch factor, R's `contr.sum` on
24/// `factor(labels)`.
25///
26/// Levels are the distinct labels in ascending order, as R orders a numeric
27/// vector's levels. The last level carries `-1` in every column.
28///
29/// ### Params
30///
31/// * `labels` - Batch label per sample
32///
33/// ### Returns
34///
35/// Row-major `n_samples * (n_levels - 1)` columns and their count, or
36/// [`EdgeErrors::InvalidArgument`] when fewer than two levels are present, the
37/// same case R's `contrasts<-` refuses.
38fn 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        // binary_search cannot fail: every label is one of the levels.
52        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
62/// Appends row-major columns onto a row-major matrix, `cbind` for this layout.
63///
64/// ### Params
65///
66/// * `left` - Row-major `n_rows * n_left`
67/// * `n_left` - Columns in `left`
68/// * `right` - Row-major `n_rows * n_right`
69/// * `n_right` - Columns in `right`
70/// * `n_rows` - Number of rows
71///
72/// ### Returns
73///
74/// Row-major `n_rows * (n_left + n_right)`.
75fn 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
84/// Checks a per-sample vector against the sample count.
85///
86/// ### Params
87///
88/// * `name` - Argument name for the error
89/// * `len` - Supplied length
90/// * `n_samples` - Required length
91///
92/// ### Returns
93///
94/// `Ok(())` or [`EdgeErrors::LengthMismatch`].
95fn 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///////////////////////
107// removeBatchEffect //
108///////////////////////
109
110/// Removes batch effects and covariates from a log-expression matrix.
111///
112/// Port of limma's `removeBatchEffect` without `group`. Fits
113/// `cbind(design, X_batch)` per gene with [`lm_fit`], where
114/// `X_batch = cbind(contr.sum(batch), contr.sum(batch2), covariates)`, and
115/// returns `x - beta_batch %*% t(X_batch)`. As in limma, covariates are
116/// column-centred first, aliased coefficients count as zero, and with no
117/// batch, second batch or covariates the input comes back unchanged.
118///
119/// ### Params
120///
121/// * `x` - Log-expression, row-major `n_genes * n_samples`. Non-finite entries
122///   are treated as missing by the fit and stay non-finite in the output.
123/// * `n_genes` - Number of genes
124/// * `n_samples` - Number of samples
125/// * `batch` - Optional batch label per sample
126/// * `batch2` - Optional second batch label per sample
127/// * `covariates` - Optional numeric covariates, row-major `n_samples * n_cov`,
128///   with `n_cov`
129/// * `design` - Optional design of interest, row-major `n_samples * n_coef`, with
130///   `n_coef`. `None` is an intercept only.
131/// * `weights` - Optional observation weights, forwarded to [`lm_fit`] as limma's
132///   `...` does
133///
134/// ### Returns
135///
136/// The corrected matrix, row-major `n_genes * n_samples`, or [`EdgeErrors`] if a
137/// shape disagrees or a batch factor has fewer than two levels.
138///
139/// ### References
140///
141/// Smyth, Statistical Applications in Genetics and Molecular Biology 3(1), 2004
142#[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, &centred, 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                // Aliased coefficients are NaN from the fit; limma zeroes them.
202                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///////////
214// Tests //
215///////////
216
217#[cfg(test)]
218mod tests {
219    // References are pasted verbatim from R's 17-digit output of
220    // `limma::removeBatchEffect` (limma 3.66.0) on the fixture below, same
221    // convention as `lm_fit`'s tests.
222    #![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    /// Absolute tolerance. References are small dyadic combinations, so anything
231    /// past round-off is a real difference.
232    const TOL: f64 = 1e-13;
233
234    /// `y[i] = ((i * 7) mod 23) / 64`, row-major 8 by 6. Dyadic, so R and Rust
235    /// see identical inputs.
236    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    /// Intercept plus a group indicator that is not orthogonal to `BATCH`.
246    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    /// `c(1, 3, 2, 5, 4, 7) / 8`.
249    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        // removeBatchEffect(y, batch = c(0, 0, 1, 1, 2, 2))
261        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        // removeBatchEffect(y, batch, design = cbind(1, c(0, 1, 1, 1, 0, 0)))
328        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        // removeBatchEffect(y, batch, batch2 = c(0, 1, 0, 1, 1, 0), design)
395        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        // removeBatchEffect(y, batch, covariates = c(1, 3, 2, 5, 4, 7) / 8, design)
462        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        // Batch identical to the design group: the batch coefficient is not
529        // estimable, limma zeroes it and hands back the input.
530        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        // removeBatchEffect(y, batch = c(7, 7, 3, 3, 7, 3)): levels sort to 3, 7
550        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        // Dyadic fixture, so the f32 cast is exact and the two agree.
649        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        // y[1, 3:4] <- NA; removeBatchEffect(y, batch = c(0, 0, 1, 1, 2, 2))
690        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        // cv <- cbind(c(1, 3, 2, 5, 4, 7) / 8, c(4, 1, 3, 0, 2, 5) / 8)
721        // removeBatchEffect(y, batch, covariates = cv, design)
722        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        // The design is only checked once there is something to remove.
810        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        // w <- matrix(1 + ((0:47 * 5) %% 7) / 4, 8, 6, byrow = TRUE)
829        // removeBatchEffect(y, batch, design, weights = w)
830        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}