_scors/
lib.rs

1mod combine;
2
3use ndarray::{Array1,ArrayView,ArrayView2,ArrayView3,ArrayViewMut1,Ix1};
4use num;
5use num::traits::float::TotalOrder;
6use numpy::{Element,PyArray,PyArray1,PyArray2,PyArray3,PyArrayDescrMethods,PyArrayDyn,PyArrayMethods,PyReadonlyArray1,PyUntypedArray,PyUntypedArrayMethods,dtype};
7use pyo3::Bound;
8use pyo3::exceptions::PyTypeError;
9use pyo3::marker::Ungil;
10use pyo3::prelude::*;
11use std::cmp::PartialOrd;
12use std::iter::{DoubleEndedIterator,repeat};
13use std::ops::AddAssign;
14
15#[derive(Clone, Copy)]
16pub enum Order {
17    ASCENDING,
18    DESCENDING
19}
20
21#[derive(Clone, Copy)]
22struct ConstWeight<F: num::Float> {
23    value: F
24}
25
26impl <F: num::Float> ConstWeight<F> {
27    fn new(value: F) -> Self {
28        return ConstWeight { value: value };
29    }
30    fn one() -> Self {
31        return Self::new(F::one());
32    }
33}
34
35pub trait Data<T: Clone>: {
36    // TODO This is necessary because it seems that there is no trait like that in rust
37    //      Maybe I am just not aware, but for now use my own trait.
38    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = T> + Clone;
39    fn get_at(&self, index: usize) -> T;
40}
41
42pub trait SortableData<T> {
43    fn argsort_unstable(&self) -> Vec<usize>;
44}
45
46impl <F: num::Float> Iterator for ConstWeight<F> {
47    type Item = F;
48    fn next(&mut self) -> Option<F> {
49        return Some(self.value);
50    }
51}
52
53impl <F: num::Float> DoubleEndedIterator for ConstWeight<F> {
54    fn next_back(&mut self) -> Option<F> {
55        return Some(self.value);
56    }
57}
58
59impl <F: num::Float> Data<F> for ConstWeight<F> {
60    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = F> + Clone {
61        return ConstWeight::new(self.value);
62    }
63
64    fn get_at(&self, _index: usize) -> F {
65        return self.value.clone();
66    }
67}
68
69impl <T: Clone> Data<T> for Vec<T> {
70    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = T> + Clone {
71        return self.iter().cloned();
72    }
73    fn get_at(&self, index: usize) -> T {
74        return self[index].clone();
75    }
76}
77
78impl SortableData<f64> for Vec<f64> {
79    fn argsort_unstable(&self) -> Vec<usize> {
80        let mut indices: Vec<usize> = (0..self.len()).collect::<Vec<_>>();
81        indices.sort_unstable_by(|i, k| self[*k].total_cmp(&self[*i]));
82        // indices.sort_unstable_by_key(|i| self[*i]);
83        return indices;
84    }
85}
86
87impl <T: Clone> Data<T> for &[T] {
88    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = T> + Clone {
89        return self.iter().cloned();
90    }
91    fn get_at(&self, index: usize) -> T {
92        return self[index].clone();
93    }
94}
95
96impl SortableData<f64> for &[f64] {
97    fn argsort_unstable(&self) -> Vec<usize> {
98        let mut indices: Vec<usize> = (0..self.len()).collect::<Vec<_>>();
99        indices.sort_unstable_by(|i, k| self[*k].total_cmp(&self[*i]));
100        return indices;
101    }
102}
103
104impl <T: Clone, const N: usize> Data<T> for [T; N] {
105    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = T> + Clone {
106        return self.iter().cloned();
107    }
108    fn get_at(&self, index: usize) -> T {
109        return self[index].clone();
110    }
111}
112
113impl <const N: usize> SortableData<f64> for [f64; N] {
114    fn argsort_unstable(&self) -> Vec<usize> {
115        let mut indices: Vec<usize> = (0..self.len()).collect::<Vec<_>>();
116        indices.sort_unstable_by(|i, k| self[*k].total_cmp(&self[*i]));
117        return indices;
118    }
119}
120
121impl <T: Clone> Data<T> for ArrayView<'_, T, Ix1> {
122    fn get_iterator(&self) -> impl DoubleEndedIterator<Item = T> + Clone {
123        return self.iter().cloned();
124    }
125    fn get_at(&self, index: usize) -> T {
126        return self[index].clone();
127    }
128}
129
130impl <F> SortableData<F> for ArrayView<'_, F, Ix1>
131where F: num::Float + TotalOrder
132{
133    fn argsort_unstable(&self) -> Vec<usize> {
134        let mut indices: Vec<usize> = (0..self.len()).collect::<Vec<_>>();
135        indices.sort_unstable_by(|i, k| self[*k].total_cmp(&self[*i]));
136        return indices;
137    }
138}
139
140pub trait BinaryLabel: Clone + Copy {
141    fn get_value(&self) -> bool;
142}
143
144impl BinaryLabel for bool {
145    fn get_value(&self) -> bool {
146        return self.clone();
147    }
148}
149
150impl BinaryLabel for u8 {
151    fn get_value(&self) -> bool {
152        return (self & 1u8) == 1u8;
153    }
154}
155
156impl BinaryLabel for u16 {
157    fn get_value(&self) -> bool {
158        return (self & 1u16) == 1u16;
159    }
160}
161
162impl BinaryLabel for u32 {
163    fn get_value(&self) -> bool {
164        return (self & 1u32) == 1u32;
165    }
166}
167
168impl BinaryLabel for u64 {
169    fn get_value(&self) -> bool {
170        return (self & 1u64) == 1u64;
171    }
172}
173
174impl BinaryLabel for i8 {
175    fn get_value(&self) -> bool {
176        return (self & 1i8) == 1i8;
177    }
178}
179
180impl BinaryLabel for i16 {
181    fn get_value(&self) -> bool {
182        return (self & 1i16) == 1i16;
183    }
184}
185
186impl BinaryLabel for i32 {
187    fn get_value(&self) -> bool {
188        return (self & 1i32) == 1i32;
189    }
190}
191
192impl BinaryLabel for i64 {
193    fn get_value(&self) -> bool {
194        return (self & 1i64) == 1i64;
195    }
196}
197
198fn select<T, I>(slice: &I, indices: &[usize]) -> Vec<T>
199where T: Copy, I: Data<T>
200{
201    let mut selection: Vec<T> = Vec::new();
202    selection.reserve_exact(indices.len());
203    for index in indices {
204        selection.push(slice.get_at(*index));
205    }
206    return selection;
207}
208
209pub trait ScoreAccumulator: num::Float + AddAssign + From<bool> + From<f32> {}
210impl<T: num::Float + AddAssign + From<bool> + From<f32>> ScoreAccumulator for T {}
211
212pub trait IntoScore<S: ScoreAccumulator>: Into<S> + num::Float {}
213impl<S: ScoreAccumulator, T: Into<S> + num::Float> IntoScore<S> for T {}
214
215
216
217pub trait ScoreSortedDescending {
218    fn _score<S: ScoreAccumulator>(&self, labels_with_weights: impl Iterator<Item = (S, (bool, S))> + Clone) -> S;
219    fn score<S, P, B, W>(&self, labels_with_weights: impl Iterator<Item = (P, (B, W))> + Clone) -> S
220    where S: ScoreAccumulator, P: IntoScore<S>, B: BinaryLabel, W: IntoScore<S>
221    {
222        return self._score(
223            labels_with_weights.map(|(p, (b, w))| -> (S, (bool, S)) { (p.into(), (b.get_value(), w.into()))})
224        )
225    }
226}
227
228
229pub fn score_sorted_iterators<S, SA, P, B, W>(
230    score: S,
231    predictions: impl Iterator<Item = P> + Clone,
232    labels: impl Iterator<Item = B> + Clone,
233    weights: impl Iterator<Item = W> + Clone,
234) -> SA
235where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel, W: IntoScore<SA> {
236    let zipped = predictions.zip(labels.zip(weights));
237    return score.score(zipped);
238}
239
240
241pub fn score_sorted_sample<S, SA, P, B, W>(
242    score: S,
243    predictions: &impl Data<P>,
244    labels: &impl Data<B>,
245    weights: &impl Data<W>,
246    order: Order,
247) -> SA
248where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel, W: IntoScore<SA> + Clone {
249    let p = predictions.get_iterator();
250    let l = labels.get_iterator();
251    let w = weights.get_iterator();
252    return match order {
253        Order::ASCENDING => score_sorted_iterators(score, p.rev(), l.rev(), w.rev()),
254        Order::DESCENDING => score_sorted_iterators(score, p, l, w),
255    };
256}
257
258
259pub fn score_maybe_sorted_sample<S, SA, P, B, W>(
260    score: S,
261    predictions: &(impl Data<P> + SortableData<P>),
262    labels: &impl Data<B>,
263    weights: Option<&impl Data<W>>,
264    order: Option<Order>,
265) -> SA
266where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel, W: IntoScore<SA> + Clone
267{
268    return match order {
269        Some(o) => {
270            match weights {
271                Some(w) => score_sorted_sample(score, predictions, labels, w, o),
272                None => score_sorted_sample(score, predictions, labels, &ConstWeight::<W>::one(), o),
273            }
274        }
275        None => {
276            let indices = predictions.argsort_unstable();
277            let sorted_labels = select(labels, &indices);
278            let sorted_predictions = select(predictions, &indices);
279            match weights {
280                Some(w) => {
281                    let sorted_weights = select(w, &indices);
282                    score_sorted_sample(score, &sorted_predictions, &sorted_labels, &sorted_weights, Order::DESCENDING)
283                }
284                None => score_sorted_sample(score, &sorted_predictions, &sorted_labels, &ConstWeight::<W>::one(), Order::DESCENDING)
285            }
286        }
287    };
288}
289
290
291pub fn score_sample<S, SA, P, B, W>(
292    score: S,
293    predictions: &(impl Data<P> + SortableData<P>),
294    labels: &impl Data<B>,
295    weights: Option<&impl Data<W>>,
296) -> SA
297
298where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel, W: IntoScore<SA> + Clone {
299    return score_maybe_sorted_sample(score, predictions, labels, weights, None);
300}
301
302
303pub fn score_two_sorted_samples<S, SA, P, B, W>(
304    score: S,
305    predictions1: impl Iterator<Item = P> + Clone,
306    label1: impl Iterator<Item = B> + Clone,
307    weight1: impl Iterator<Item = W> + Clone,
308    predictions2: impl Iterator<Item = P> + Clone,
309    label2: impl Iterator<Item = B> + Clone,
310    weight2: impl Iterator<Item = W> + Clone,
311) -> SA
312where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel + PartialOrd, W: IntoScore<SA>
313{
314    return score_two_sorted_samples_zipped(
315        score,
316        predictions1.zip(label1.zip(weight1)),
317        predictions2.zip(label2.zip(weight2)),
318    );
319}
320
321
322pub fn score_two_sorted_samples_zipped<S, SA, P, B, W>(
323    score: S,
324    iter1: impl Iterator<Item = (P, (B, W))> + Clone,
325    iter2: impl Iterator<Item = (P, (B, W))> + Clone,
326) -> SA
327where S: ScoreSortedDescending, SA: ScoreAccumulator, P: IntoScore<SA>, B: BinaryLabel + PartialOrd, W: IntoScore<SA>
328{
329    let combined_iter = combine::combine::CombineIterDescending::new(iter1, iter2);
330    return score.score(combined_iter);
331}
332
333
334struct AveragePrecision {
335    
336}
337
338
339impl AveragePrecision {
340    fn new() -> Self {
341        return AveragePrecision{};
342    }
343}
344
345
346#[derive(Clone,Copy,Debug)]
347struct Positives<P>
348where P: num::Float + From<bool> + AddAssign
349{
350    tps: P,
351    fps: P,
352}
353
354impl <P> Positives<P>
355where P: num::Float + From<bool> + AddAssign
356{
357    fn new(tps: P, fps: P) -> Self {
358        return Positives { tps, fps };
359    }
360
361    fn zero() -> Self {
362        return Positives::new(P::zero(), P::zero());
363    }
364
365    fn add(&mut self, label: bool, weight: P) {
366        let label: P = label.into();
367        let tp = weight * label;
368        let fp = weight - tp;  // (weight*(1 -label) = weight - weight * label = weight - tp)
369        self.tps += tp;
370        self.fps += fp;
371    }
372
373    fn positives_sum(&self) -> P {
374        return self.tps + self.fps;
375    }
376
377    fn precision(&self) -> P {
378        return self.tps / self.positives_sum();
379    }
380}
381
382
383impl ScoreSortedDescending for AveragePrecision {
384    fn _score<S: ScoreAccumulator>(&self, mut labels_with_weights: impl Iterator<Item = (S, (bool, S))> + Clone) -> S
385    {
386        let mut positives: Positives<S> = Positives::zero();
387        let mut last_p: S = f32::NAN.into();
388        let mut last_tps: S = S::zero();
389        let mut ap: S = S::zero();
390
391        // TODO can we unify this preparation step with the loop?
392        match labels_with_weights.next() {
393            None => (), // TODO: Sohuld we return an error in this case?
394            Some((p, (label, w))) => {
395                positives.add(label, w);
396                last_p = p;
397            }
398        }
399        
400        for (p, (label, w)) in labels_with_weights {
401            if last_p != p {
402                ap += (positives.tps - last_tps) * positives.precision();
403                last_p = p;
404                last_tps = positives.tps;
405            }
406            positives.add(label.get_value(), w.into());
407        }
408
409        ap += (positives.tps - last_tps) * positives.precision();
410        
411        // Special case for tps == 0 following sklearn
412        // https://github.com/scikit-learn/scikit-learn/blob/5cce87176a530d2abea45b5a7e5a4d837c481749/sklearn/metrics/_ranking.py#L1032-L1039
413        // I.e. if tps is 0.0, there are no positive samples in labels: Either all labels are 0, or all weights (for positive labels) are 0
414        return if positives.tps == S::zero() {
415            S::zero()
416        } else {
417            ap / positives.tps
418        };
419    }
420}
421
422
423struct RocAuc {
424
425}
426
427
428impl RocAuc {
429    fn new() -> Self {
430        return RocAuc { };
431    }
432}
433
434
435impl ScoreSortedDescending for RocAuc {
436    fn _score<S: ScoreAccumulator>(&self, mut labels_with_weights: impl Iterator<Item = (S, (bool, S))> + Clone) -> S
437    {
438        let mut positives: Positives<S> = Positives::zero();
439        let mut last_p: S = f32::NAN.into();
440        let mut last_counted_fp = S::zero();
441        let mut last_counted_tp = S::zero();
442        let mut area_under_curve = S::zero();
443
444        // TODO can we unify this preparation step with the loop?
445        match labels_with_weights.next() {
446            None => (), // TODO: Should we return an error in this case?
447            Some((p, (label, w))) => {
448                positives.add(label, w);
449                last_p = p;
450            }
451        }
452        
453        for (p, (label, w)) in labels_with_weights {
454            if last_p != p {
455                area_under_curve += area_under_line_segment(
456                    last_counted_fp,
457                    positives.fps,
458                    last_counted_tp,
459                    positives.tps,
460                );
461                last_counted_fp = positives.fps;
462                last_counted_tp = positives.tps;
463                last_p = p;
464            }
465            positives.add(label, w);
466        }
467        area_under_curve += area_under_line_segment(
468            last_counted_fp,
469            positives.fps,
470            last_counted_tp,
471            positives.tps,
472        );
473        return area_under_curve / (positives.tps * positives.fps);
474    }
475}
476
477
478struct RocAucWithMaxFPR {
479    max_fpr: f32,
480}
481
482
483impl RocAucWithMaxFPR {
484    fn new(max_fpr: f32) -> Self {
485        return RocAucWithMaxFPR { max_fpr };
486    }
487
488    fn get_positive_sum<B, W>(labels_with_weights: impl Iterator<Item = (B, W)>) -> Positives<W>
489    where B: BinaryLabel, W: num::Float + From::<bool> + AddAssign
490    {
491        let mut positives: Positives<W>  = Positives::zero();
492        for (label, weight) in labels_with_weights {
493            positives.add(label.get_value(), weight);
494        }
495        return positives;
496    }
497}
498
499
500impl ScoreSortedDescending for RocAucWithMaxFPR {
501    fn _score<S: ScoreAccumulator>(&self, mut labels_with_weights: impl Iterator<Item = (S, (bool, S))> + Clone) -> S
502    {
503        let total_positives = Self::get_positive_sum(labels_with_weights.clone().map(|(_a, b)| b));
504        let max_fpr: S = self.max_fpr.into();
505        let false_positive_cutoff = max_fpr * total_positives.fps;
506
507        let mut positives: Positives<S> = Positives::zero();
508        let mut last_p: S = f32::NAN.into();
509        let mut last_counted_fp = S::zero();
510        let mut last_counted_tp = S::zero();
511        let mut area_under_curve = S::zero();
512
513        // TODO can we unify this preparation step with the loop?
514        match labels_with_weights.next() {
515            None => (), // TODO: Should we return an error in this case?
516            Some((p, (label, w))) => {
517                positives.add(label, w);
518                last_p = p;
519            }
520        }
521        
522        for (p, (label, w)) in labels_with_weights {
523            if last_p != p {
524                area_under_curve += area_under_line_segment(
525                    last_counted_fp,
526                    positives.fps,
527                    last_counted_tp,
528                    positives.tps,
529                );
530                last_counted_fp = positives.fps;
531                last_counted_tp = positives.tps;
532                last_p = p;
533            }
534            let mut next_pos = positives.clone();
535            next_pos.add(label, w);
536            if next_pos.fps > false_positive_cutoff {
537                let dx = next_pos.fps - positives.fps;
538                let dy = next_pos.tps - positives.tps;
539                positives = Positives::new(
540                    positives.tps + dy * false_positive_cutoff / dx,
541                    false_positive_cutoff,
542                );
543                break;
544            }
545            else {
546                positives = next_pos;
547            }
548        }
549
550        area_under_curve += area_under_line_segment(
551            last_counted_fp,
552            positives.fps,
553            last_counted_tp,
554            positives.tps,
555        );
556        
557        let normalized_area_under_curve = area_under_curve / (total_positives.tps * total_positives.fps);
558        let one_half: S = 0.5f32.into(); 
559        let min_area = one_half * max_fpr * max_fpr;
560        let max_area = max_fpr;
561        return one_half * (S::one() + (normalized_area_under_curve - min_area) / (max_area - min_area));
562    }
563}
564
565
566struct RocAucWithOptionalMaxFPR {
567    // TODO: Can we have a single implementation for this and RocAuc?
568    //       This would add an unncessary check to RocAuc but performance
569    //       penalty may be negligible.
570    max_fpr: Option<f32>,
571}
572
573impl RocAucWithOptionalMaxFPR {
574    fn new(max_fpr: Option<f32>) -> Self {
575        return Self { max_fpr };
576    }
577}
578
579
580impl ScoreSortedDescending for RocAucWithOptionalMaxFPR {
581    fn _score<S: ScoreAccumulator>(&self, labels_with_weights: impl Iterator<Item = (S, (bool, S))> + Clone) -> S
582    {
583        return match self.max_fpr {
584            Some(mfpr) => RocAucWithMaxFPR::new(mfpr).score(labels_with_weights),
585            None => RocAuc::new().score(labels_with_weights),
586        }
587    }
588}
589
590
591pub fn average_precision<S, P, B, W>(
592    predictions: &(impl Data<P> + SortableData<P>),
593    labels: &impl Data<B>,
594    weights: Option<&impl Data<W>>,
595    order: Option<Order>,
596) -> S
597where S: ScoreAccumulator, P: IntoScore<S>, B: BinaryLabel, W: IntoScore<S> + Clone
598{
599    return score_maybe_sorted_sample(AveragePrecision::new(), predictions, labels, weights, order);
600}
601
602
603pub fn roc_auc<S, P, B, W>(
604    predictions: &(impl Data<P> + SortableData<P>),
605    labels: &impl Data<B>,
606    weights: Option<&impl Data<W>>,
607    order: Option<Order>,
608    max_fpr: Option<f32>,
609) -> S
610where S: ScoreAccumulator, P: IntoScore<S>, B: BinaryLabel, W: IntoScore<S> + Clone
611{
612    return score_maybe_sorted_sample(RocAucWithOptionalMaxFPR::new(max_fpr), predictions, labels, weights, order);
613}
614
615
616fn area_under_line_segment<P>(x0: P, x1: P, y0: P, y1: P) -> P
617where P: num::Float + From<f32>
618{
619    let dx = x1 - x0;
620    let dy = y1 - y0;
621    let one_half: P = 0.5f32.into();
622    return dx * y0 + dy * dx * one_half;
623}
624
625
626/// Accumulate `row` into `sum` element-wise.
627///
628/// When monomorphized with `Copied<slice::Iter>` LLVM proves unit stride and
629/// emits SIMD.  When called with ndarray's strided `Iter` it stays scalar.
630#[inline]
631fn accum_row<F: num::Float + AddAssign>(row: impl Iterator<Item = F>, sum: &mut [F]) {
632    for (f, s) in row.zip(sum.iter_mut()) {
633        *s += f;
634    }
635}
636
637/// Compute (prod_sum, m_sqs, l_sqs) for one replicate row.
638///
639/// Same vectorization note as `accum_row`.
640#[inline]
641fn score_row<F: num::Float + AddAssign>(
642    row: impl Iterator<Item = F>,
643    sum: impl Iterator<Item = F>,
644    loo_weight_factor: F,
645) -> (F, F, F) {
646    let mut prod_sum = F::zero();
647    let mut m_sqs = F::zero();
648    let mut l_sqs = F::zero();
649    for (f, s) in row.zip(sum) {
650        let m_f = f;
651        let l_f = (s - f) * loo_weight_factor;
652        prod_sum += m_f * l_f;
653        m_sqs += m_f * m_f;
654        l_sqs += l_f * l_f;
655    }
656    (prod_sum, m_sqs, l_sqs)
657}
658
659/// Core two-pass loop shared by the contiguous and strided paths.
660///
661/// `rows` must be `Clone` because the algorithm makes two passes: one to
662/// accumulate `sum`, one to score.  The `Clone` is free — the iterator is
663/// just a pointer and two `usize`s on the stack.
664///
665/// LLVM monomorphizes separately for each `Row` type: when called with
666/// `Chunks`-derived slice iterators it emits SIMD; when called with
667/// ndarray's strided iterators it stays scalar.
668///
669/// Note: `impl Fn(&ArrayView1) -> I` would be cleaner at the call site but
670/// hits a Rust lifetime limitation — `as_slice()` ties its `&[F]` return to
671/// `&self` (the view reference), not to the underlying data lifetime, so the
672/// returned iterator borrows a local and the compiler rejects it.  The
673/// iterator solution avoids this: `chunks()` borrows directly from `mat`
674/// and `into_iter()` transfers the data lifetime from the view.
675#[inline]
676fn loo_cossim_loops<F, Row>(
677    rows: impl Iterator<Item = Row> + Clone,
678    sum: &mut [F],
679    loo_weight_factor: F,
680    num_replicates: usize,
681) -> F
682where
683    F: num::Float + AddAssign,
684    Row: Iterator<Item = F>,
685{
686    for row in rows.clone() {
687        accum_row(row, sum);
688    }
689    let mut result = F::zero();
690    for row in rows {
691        let (prod_sum, m_sqs, l_sqs) = score_row(row, sum.iter().copied(), loo_weight_factor);
692        result += prod_sum / (m_sqs * l_sqs).sqrt();
693    }
694    result / F::from(num_replicates).unwrap()
695}
696
697pub fn loo_cossim<F: num::Float + AddAssign>(mat: &ArrayView2<'_, F>, replicate_sum: &mut ArrayViewMut1<'_, F>) -> F {
698    let num_replicates = mat.shape()[0];
699    let loo_weight_factor = F::from(1).unwrap() / F::from(num_replicates - 1).unwrap();
700    let ncols = mat.shape()[1];
701    // replicate_sum is always created from Array1::zeros() inside Rust, so
702    // it is guaranteed to be contiguous.
703    let sum = replicate_sum.as_slice_mut().unwrap();
704
705    // mat.as_slice() succeeds iff the matrix is fully C-contiguous (strides
706    // == [ncols, 1]).  chunks() then yields &[F] slices that borrow directly
707    // from mat's data so LLVM sees plain pointer arithmetic and emits SIMD.
708    // For non-contiguous input (strided or Fortran-order) outer_iter() +
709    // into_iter() transfer the data lifetime from the view correctly.
710    //
711    // In the primary call path (loo_cossim_many, univariate sampling) each
712    // 2-D slice has strides (ncols, 1) and as_slice() always succeeds.
713    if let Some(mat_slice) = mat.as_slice() {
714        loo_cossim_loops(mat_slice.chunks(ncols).map(|s| s.iter().copied()), sum, loo_weight_factor, num_replicates)
715    } else {
716        loo_cossim_loops(mat.outer_iter().map(|r| r.into_iter().copied()), sum, loo_weight_factor, num_replicates)
717    }
718}
719
720
721pub fn loo_cossim_single<F: num::Float + AddAssign>(mat: &ArrayView2<'_, F>) -> F {
722    let mut replicate_sum = Array1::<F>::zeros(mat.shape()[1]);
723    return loo_cossim(mat, &mut replicate_sum.view_mut());
724}
725
726
727pub fn loo_cossim_many<F: num::Float + AddAssign>(mat: &ArrayView3<'_, F>) -> Array1<F> {
728    let mut cossims = Array1::<F>::zeros(mat.shape()[0]);
729    let mut replicate_sum = Array1::<F>::zeros(mat.shape()[2]);
730    for (m, c) in mat.outer_iter().zip(cossims.iter_mut()) {
731        replicate_sum.fill(F::zero());
732        *c = loo_cossim(&m, &mut replicate_sum.view_mut());
733    }
734    return cossims;
735}
736
737
738// Python bindings
739#[pyclass(eq, eq_int, from_py_object, name="Order")]
740#[derive(Clone, Copy, PartialEq)]
741pub enum PyOrder {
742    ASCENDING,
743    DESCENDING
744}
745
746fn py_order_as_order(order: PyOrder) -> Order {
747    return match order {
748        PyOrder::ASCENDING => Order::ASCENDING,
749        PyOrder::DESCENDING => Order::DESCENDING,
750    }
751}
752
753trait PyScoreGeneric<S: ScoreSortedDescending>: Ungil + Sync {
754
755    fn get_score(&self) -> S;
756
757    fn score_py<'py, P, B, W>(
758        &self,
759        py: Python<'py>,
760        labels: PyReadonlyArray1<'py, B>,
761        predictions: PyReadonlyArray1<'py, P>,
762        weights: Option<PyReadonlyArray1<'py, W>>,
763        order: Option<PyOrder>,
764    ) -> P
765    where P: ScoreAccumulator + Element + TotalOrder, B: BinaryLabel + Element, W: IntoScore<P> + Element
766    {
767        let labels = labels.as_array();
768        let predictions = predictions.as_array();
769        let order = order.map(py_order_as_order);
770        let score = match weights {
771            Some(weight) => {
772                let w = weight.as_array();
773                py.detach(move || {
774                    score_maybe_sorted_sample(self.get_score(), &predictions, &labels, Some(&w), order)
775                })
776            },
777            None => py.detach(move || {
778                score_maybe_sorted_sample(self.get_score(), &predictions, &labels, None::<&Vec<W>>, order)
779            })
780        };
781        return score;
782    }
783
784    fn score_two_sorted_samples_py_generic<'py, B, F, W, B1, B2, F1, F2, W1, W2>(
785        &self,
786        py: Python<'py>,
787        labels1: PyReadonlyArray1<'py, B1>,
788        predictions1: PyReadonlyArray1<'py, F1>,
789        weights1: Option<PyReadonlyArray1<'py, W1>>,
790        labels2: PyReadonlyArray1<'py, B1>,
791        predictions2: PyReadonlyArray1<'py, F2>,
792        weights2: Option<PyReadonlyArray1<'py, W2>>,
793    ) -> F
794    where B: BinaryLabel + PartialOrd, F: ScoreAccumulator + TotalOrder + Ungil, W: IntoScore<F>, B1: Element + Into<B> + Clone, B2: Element + Into<B> + Clone, F1: Element + Into<F> + Clone, F2: Element + Into<F> + Clone, W1: Element + Into<W> + Clone, W2: Element + Into<W> + Clone
795    {
796        let l1 = labels1.as_array().into_iter().cloned().map(|l| -> B { l.into() });
797        let l2 = labels2.as_array().into_iter().cloned().map(|l| -> B { l.into() });
798        let p1 = predictions1.as_array().into_iter().cloned().map(|f| -> F { f.into() });
799        let p2 = predictions2.as_array().into_iter().cloned().map(|f| -> F { f.into() });
800
801
802        return match (weights1, weights2) {
803            (None, None) => {
804                py.detach(move || {
805                    score_two_sorted_samples(self.get_score(), p1, l1, repeat(W::one()), p2, l2, repeat(W::one()))
806                })
807            }
808            (Some(w1), None) => {
809                let w1i = w1.as_array().into_iter().cloned().map(|w| -> W { w.into() });
810                py.detach(move || {
811                    score_two_sorted_samples(self.get_score(), p1, l1, w1i, p2, l2, repeat(W::one()))
812                })
813            }
814            (None, Some(w2)) => {
815                let w2i = w2.as_array().into_iter().cloned().map(|w| -> W { w.into() });
816                py.detach(move || {
817                    score_two_sorted_samples(self.get_score(), p1, l1, repeat(W::one()), p2, l2, w2i)
818                })
819            }
820            (Some(w1), Some(w2)) =>  {
821                let w1i = w1.as_array().into_iter().cloned().map(|w| -> W { w.into() });
822                let w2i = w2.as_array().into_iter().cloned().map(|w| -> W { w.into() });
823                py.detach(move || {
824                    score_two_sorted_samples(self.get_score(), p1, l1, w1i, p2, l2, w2i)
825                })
826            }
827        };
828    }
829}
830
831struct AveragePrecisionPyGeneric {
832
833}
834
835impl AveragePrecisionPyGeneric {
836    fn new() -> Self {
837        return AveragePrecisionPyGeneric {};
838    }
839}
840
841impl PyScoreGeneric<AveragePrecision> for AveragePrecisionPyGeneric {
842    fn get_score(&self) -> AveragePrecision {
843        return AveragePrecision::new();
844    }
845}
846
847struct RocAucPyGeneric {
848    max_fpr: Option<f32>,
849}
850
851impl RocAucPyGeneric {
852    fn new(max_fpr: Option<f32>) -> Self {
853        return RocAucPyGeneric { max_fpr: max_fpr };
854    }
855}
856
857impl PyScoreGeneric<RocAucWithOptionalMaxFPR> for RocAucPyGeneric {
858    fn get_score(&self) -> RocAucWithOptionalMaxFPR {
859        return RocAucWithOptionalMaxFPR::new(self.max_fpr);
860    }
861}
862
863// https://stackoverflow.com/questions/70128978/how-to-define-different-function-names-with-a-macro
864// https://stackoverflow.com/questions/70872059/using-a-rust-macro-to-generate-a-function-with-variable-parameters
865// https://doc.rust-lang.org/rust-by-example/macros/designators.html
866// https://users.rust-lang.org/t/is-there-a-way-to-convert-given-identifier-to-a-string-in-a-macro/42907
867macro_rules! average_precision_py {
868    ($fname: ident, $pyname:literal, $label_type:ty, $prediction_type:ty, $weight_type:ty) => {
869        #[pyfunction(name = $pyname)]
870        #[pyo3(signature = (labels, predictions, *, weights=None, order=None))]
871        pub fn $fname<'py>(
872            py: Python<'py>,
873            labels: PyReadonlyArray1<'py, $label_type>,
874            predictions: PyReadonlyArray1<'py, $prediction_type>,
875            weights: Option<PyReadonlyArray1<'py, $weight_type>>,
876            order: Option<PyOrder>
877        ) -> $prediction_type
878        {
879            return AveragePrecisionPyGeneric::new().score_py(py, labels, predictions, weights, order);
880        }
881    };
882    ($fname: ident, $pyname:literal, $label_type:ty, $prediction_type:ty, $weight_type:ty, $py_module:ident) => {
883        average_precision_py!($fname, $pyname, $label_type, $prediction_type, $weight_type);
884        $py_module.add_function(wrap_pyfunction!($fname, $py_module)?).unwrap();
885    };
886}
887
888
889macro_rules! roc_auc_py {
890    ($fname: ident, $pyname:literal, $label_type:ty, $prediction_type:ty, $weight_type:ty) => {
891        #[pyfunction(name = $pyname)]
892        #[pyo3(signature = (labels, predictions, *, weights=None, order=None, max_fpr=None))]
893        pub fn $fname<'py>(
894            py: Python<'py>,
895            labels: PyReadonlyArray1<'py, $label_type>,
896            predictions: PyReadonlyArray1<'py, $prediction_type>,
897            weights: Option<PyReadonlyArray1<'py, $weight_type>>,
898            order: Option<PyOrder>,
899            max_fpr: Option<f32>,
900        ) -> $prediction_type
901        {
902            return RocAucPyGeneric::new(max_fpr).score_py(py, labels, predictions, weights, order);
903        }
904    };
905    ($fname: ident, $pyname:literal, $label_type:ty, $prediction_type:ty, $weight_type: ty, $py_module:ident) => {
906        roc_auc_py!($fname, $pyname, $label_type, $prediction_type, $weight_type);
907        $py_module.add_function(wrap_pyfunction!($fname, $py_module)?).unwrap();
908    };
909}
910
911
912macro_rules! average_precision_on_two_sorted_samples_py {
913    ($fname: ident, $pyname:literal, $lt:ty, $pt:ty, $wt:ty, $lt1:ty, $pt1:ty, $wt1:ty, $lt2:ty, $pt2:ty, $wt2: ty) => {
914        #[pyfunction(name = $pyname)]
915        #[pyo3(signature = (labels1, predictions1, weights1, labels2, predictions2, weights2, *))]
916        pub fn $fname<'py>(
917            py: Python<'py>,
918            labels1: PyReadonlyArray1<'py, $lt1>,
919            predictions1: PyReadonlyArray1<'py, $pt1>,
920            weights1: Option<PyReadonlyArray1<'py, $wt1>>,
921            labels2: PyReadonlyArray1<'py, $lt2>,
922            predictions2: PyReadonlyArray1<'py, $pt2>,
923            weights2: Option<PyReadonlyArray1<'py, $wt2>>,
924        ) -> $pt
925        {
926            return AveragePrecisionPyGeneric::new().score_two_sorted_samples_py_generic::<$lt, $pt, $wt, $lt1, $lt2, $pt1, $pt2, $wt1, $wt2>(py, labels1, predictions1, weights1, labels2, predictions2, weights2);
927        }
928    };
929    ($fname: ident, $pyname:literal, $lt:ty, $pt:ty, $wt:ty, $lt1:ty, $pt1:ty, $wt1:ty, $lt2:ty, $pt2:ty, $wt2: ty, $py_module:ident) => {
930        average_precision_on_two_sorted_samples_py!($fname, $pyname, $lt, $pt, $wt, $lt1, $pt1, $wt1, $lt2, $pt2, $wt2);
931        $py_module.add_function(wrap_pyfunction!($fname, $py_module)?).unwrap();
932    };
933}
934
935
936macro_rules! roc_auc_on_two_sorted_samples_py {
937    ($fname: ident, $pyname:literal, $lt:ty, $pt:ty, $wt:ty, $lt1:ty, $pt1:ty, $wt1:ty, $lt2:ty, $pt2:ty, $wt2: ty) => {
938        #[pyfunction(name = $pyname)]
939        #[pyo3(signature = (labels1, predictions1, weights1, labels2, predictions2, weights2, *, max_fpr=None))]
940        pub fn $fname<'py>(
941            py: Python<'py>,
942            labels1: PyReadonlyArray1<'py, $lt1>,
943            predictions1: PyReadonlyArray1<'py, $pt1>,
944            weights1: Option<PyReadonlyArray1<'py, $wt1>>,
945            labels2: PyReadonlyArray1<'py, $lt2>,
946            predictions2: PyReadonlyArray1<'py, $pt2>,
947            weights2: Option<PyReadonlyArray1<'py, $wt2>>,
948            max_fpr: Option<f32>,
949        ) -> $pt
950        {
951            return RocAucPyGeneric::new(max_fpr).score_two_sorted_samples_py_generic::<$lt, $pt, $wt, $lt1, $lt2, $pt1, $pt2, $wt1, $wt2>(py, labels1, predictions1, weights1, labels2, predictions2, weights2);
952        }
953    };
954    ($fname: ident, $pyname:literal, $lt:ty, $pt:ty, $wt:ty, $lt1:ty, $pt1:ty, $wt1:ty, $lt2:ty, $pt2:ty, $wt2: ty, $py_module:ident) => {
955        roc_auc_on_two_sorted_samples_py!($fname, $pyname, $lt, $pt, $wt, $lt1, $pt1, $wt1, $lt2, $pt2, $wt2);
956        $py_module.add_function(wrap_pyfunction!($fname, $py_module)?).unwrap();
957    };
958}
959
960
961#[pyfunction(name = "loo_cossim")]
962#[pyo3(signature = (data))]
963pub fn loo_cossim_py<'py>(
964    py: Python<'py>,
965    data: &Bound<'py, PyUntypedArray>
966) -> PyResult<f64> {
967    if data.ndim() != 2 {
968        return Err(PyTypeError::new_err(format!("Expected 2-dimensional array for data (samples x features) but found {} dimenisons.", data.ndim())));
969    }
970
971    let dt = data.dtype();
972    if dt.is_equiv_to(&dtype::<f32>(py)) {
973        let typed_data = data.cast::<PyArray2<f32>>().unwrap().readonly();
974        let array = typed_data.as_array();
975        let score = py.detach(move || {
976            loo_cossim_single(&array)
977        });
978        return Ok(score as f64);
979    }
980    if dt.is_equiv_to(&dtype::<f64>(py)) {
981        let typed_data = data.cast::<PyArray2<f64>>().unwrap().readonly();
982        let array = typed_data.as_array();
983        let score = py.detach(move || {
984            loo_cossim_single(&array)
985        });
986        return Ok(score);
987    }
988    return Err(PyTypeError::new_err(format!("Only float32 and float64 data supported, but found {}", dt)));
989}
990
991pub fn loo_cossim_many_generic_py<'py, F: num::Float + AddAssign + Element>(
992    py: Python<'py>,
993    data: &Bound<'py, PyArrayDyn<F>>
994) -> PyResult<Bound<'py, PyArray1<F>>> {
995    if data.ndim() != 3 {
996        return Err(PyTypeError::new_err(format!("Expected 3-dimensional array for data (outer(?) x samples x features) but found {} dimenisons.", data.ndim())));
997    }
998    let typed_data = data.cast::<PyArray3<F>>().unwrap().readonly();
999    let array = typed_data.as_array();
1000    let score = py.detach(move || {
1001        loo_cossim_many(&array)
1002    });
1003    // TODO how can we return this generically without making a copy at the end?
1004    let score_py = PyArray::from_owned_array(py, score);
1005    return Ok(score_py);
1006}
1007
1008#[pyfunction(name = "loo_cossim_many_f64")]
1009#[pyo3(signature = (data))]
1010pub fn loo_cossim_many_py_f64<'py>(
1011    py: Python<'py>,
1012    data: &Bound<'py, PyUntypedArray>
1013) -> PyResult<Bound<'py, PyArray1<f64>>> {
1014    if data.ndim() != 3 {
1015        return Err(PyTypeError::new_err(format!("Expected 3-dimensional array for data (outer(?) x samples x features) but found {} dimenisons.", data.ndim())));
1016    }
1017
1018    let dt = data.dtype();
1019    if !dt.is_equiv_to(&dtype::<f64>(py)) {
1020        return Err(PyTypeError::new_err(format!("Only float64 data supported, but found {}", dt)));
1021    }
1022    let typed_data = data.cast::<PyArrayDyn<f64>>().unwrap();
1023    return loo_cossim_many_generic_py(py, typed_data);
1024}
1025
1026#[pyfunction(name = "loo_cossim_many_f32")]
1027#[pyo3(signature = (data))]
1028pub fn loo_cossim_many_py_f32<'py>(
1029    py: Python<'py>,
1030    data: &Bound<'py, PyUntypedArray>
1031) -> PyResult<Bound<'py, PyArray1<f32>>> {
1032    if data.ndim() != 3 {
1033        return Err(PyTypeError::new_err(format!("Expected 3-dimensional array for data (outer(?) x samples x features) but found {} dimenisons.", data.ndim())));
1034    }
1035
1036    let dt = data.dtype();
1037    if !dt.is_equiv_to(&dtype::<f32>(py)) {
1038        return Err(PyTypeError::new_err(format!("Only float32 data supported, but found {}", dt)));
1039    }
1040    let typed_data = data.cast::<PyArrayDyn<f32>>().unwrap();
1041    return loo_cossim_many_generic_py(py, typed_data);
1042}
1043
1044#[pymodule(name = "_scors")]
1045fn scors(m: &Bound<'_, PyModule>) -> PyResult<()> {
1046    average_precision_py!(average_precision_bool_f32, "average_precision_bool_f32", bool, f32, f32, m);
1047    average_precision_py!(average_precision_i8_f32, "average_precision_i8_f32", i8, f32, f32, m);
1048    average_precision_py!(average_precision_i16_f32, "average_precision_i16_f32", i16, f32, f32, m);
1049    average_precision_py!(average_precision_i32_f32, "average_precision_i32_f32", i32, f32, f32, m);
1050    average_precision_py!(average_precision_i64_f32, "average_precision_i64_f32", i64, f32, f32, m);
1051    average_precision_py!(average_precision_u8_f32, "average_precision_u8_f32", u8, f32, f32, m);
1052    average_precision_py!(average_precision_u16_f32, "average_precision_u16_f32", u16, f32, f32, m);
1053    average_precision_py!(average_precision_u32_f32, "average_precision_u32_f32", u32, f32, f32, m);
1054    average_precision_py!(average_precision_u64_f32, "average_precision_u64_f32", u64, f32, f32, m);
1055    average_precision_py!(average_precision_bool_f64, "average_precision_bool_f64", bool, f64, f64, m);
1056    average_precision_py!(average_precision_i8_f64, "average_precision_i8_f64", i8, f64, f64, m);
1057    average_precision_py!(average_precision_i16_f64, "average_precision_i16_f64", i16, f64, f64, m);
1058    average_precision_py!(average_precision_i32_f64, "average_precision_i32_f64", i32, f64, f64, m);
1059    average_precision_py!(average_precision_i64_f64, "average_precision_i64_f64", i64, f64, f64, m);
1060    average_precision_py!(average_precision_u8_f64, "average_precision_u8_f64", u8, f64, f64, m);
1061    average_precision_py!(average_precision_u16_f64, "average_precision_u16_f64", u16, f64, f64, m);
1062    average_precision_py!(average_precision_u32_f64, "average_precision_u32_f64", u32, f64, f64, m);
1063    average_precision_py!(average_precision_u64_f64, "average_precision_u64_f64", u64, f64, f64, m);
1064
1065    roc_auc_py!(roc_auc_bool_f32, "roc_auc_bool_f32", bool, f32, f32, m);
1066    roc_auc_py!(roc_auc_i8_f32, "roc_auc_i8_f32", i8, f32, f32, m);
1067    roc_auc_py!(roc_auc_i16_f32, "roc_auc_i16_f32", i16, f32, f32, m);
1068    roc_auc_py!(roc_auc_i32_f32, "roc_auc_i32_f32", i32, f32, f32, m);
1069    roc_auc_py!(roc_auc_i64_f32, "roc_auc_i64_f32", i64, f32, f32, m);
1070    roc_auc_py!(roc_auc_u8_f32, "roc_auc_u8_f32", u8, f32, f32, m);
1071    roc_auc_py!(roc_auc_u16_f32, "roc_auc_u16_f32", u16, f32, f32, m);
1072    roc_auc_py!(roc_auc_u32_f32, "roc_auc_u32_f32", u32, f32, f32, m);
1073    roc_auc_py!(roc_auc_u64_f32, "roc_auc_u64_f32", u64, f32, f32, m);
1074    roc_auc_py!(roc_auc_bool_f64, "roc_auc_bool_f64", bool, f64, f64, m);
1075    roc_auc_py!(roc_auc_i8_f64, "roc_auc_i8_f64", i8, f64, f64, m);
1076    roc_auc_py!(roc_auc_i16_f64, "roc_auc_i16_f64", i16, f64, f64, m);
1077    roc_auc_py!(roc_auc_i32_f64, "roc_auc_i32_f64", i32, f64, f64, m);
1078    roc_auc_py!(roc_auc_i64_f64, "roc_auc_i64_f64", i64, f64, f64, m);
1079    roc_auc_py!(roc_auc_u8_f64, "roc_auc_u8_f64", u8, f64, f64, m);
1080    roc_auc_py!(roc_auc_u16_f64, "roc_auc_u16_f64", u16, f64, f64, m);
1081    roc_auc_py!(roc_auc_u32_f64, "roc_auc_u32_f64", u32, f64, f64, m);
1082    roc_auc_py!(roc_auc_u64_f64, "roc_auc_u64_f64", u64, f64, f64, m);
1083
1084    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_bool_f32, "average_precision_on_two_sorted_samples_bool_f32", bool, f32, f32, bool, f32, f32, bool, f32, f32, m);
1085    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i8_f32, "average_precision_on_two_sorted_samples_i8_f32", i8, f32, f32, i8, f32, f32, i8, f32, f32, m);
1086    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i16_f32, "average_precision_on_two_sorted_samples_i16_f32", i16, f32, f32, i16, f32, f32, i16, f32, f32, m);
1087    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i32_f32, "average_precision_on_two_sorted_samples_i32_f32", i32, f32, f32, i32, f32, f32, i32, f32, f32, m);
1088    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i64_f32, "average_precision_on_two_sorted_samples_i64_f32", i64, f32, f32, i64, f32, f32, i64, f32, f32, m);
1089    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u8_f32, "average_precision_on_two_sorted_samples_u8_f32", u8, f32, f32, u8, f32, f32, u8, f32, f32, m);
1090    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u16_f32, "average_precision_on_two_sorted_samples_u16_f32", u16, f32, f32, u16, f32, f32, u16, f32, f32, m);
1091    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u32_f32, "average_precision_on_two_sorted_samples_u32_f32", u32, f32, f32, u32, f32, f32, u32, f32, f32, m);
1092    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u64_f32, "average_precision_on_two_sorted_samples_u64_f32", u64, f32, f32, u64, f32, f32, u64, f32, f32, m);
1093    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_bool_f64, "average_precision_on_two_sorted_samples_bool_f64", bool, f64, f64, bool, f64, f64, bool, f64, f64, m);
1094    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i8_f64, "average_precision_on_two_sorted_samples_i8_f64", i8, f64, f64, i8, f64, f64, i8, f64, f64, m);
1095    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i16_f64, "average_precision_on_two_sorted_samples_i16_f64", i16, f64, f64, i16, f64, f64, i16, f64, f64, m);
1096    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i32_f64, "average_precision_on_two_sorted_samples_i32_f64", i32, f64, f64, i16, f64, f64, i16, f64, f64, m);
1097    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_i64_f64, "average_precision_on_two_sorted_samples_i64_f64", i64, f64, f64, i64, f64, f64, i64, f64, f64, m);
1098    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u8_f64, "average_precision_on_two_sorted_samples_u8_f64", u8, f64, f64, u8, f64, f64, u8, f64, f64, m);
1099    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u16_f64, "average_precision_on_two_sorted_samples_u16_f64", u16, f64, f64, u16, f64, f64, u16, f64, f64, m);
1100    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u32_f64, "average_precision_on_two_sorted_samples_u32_f64", u32, f64, f64, u32, f64, f64, u32, f64, f64, m);
1101    average_precision_on_two_sorted_samples_py!(average_precision_on_two_sorted_samples_u64_f64, "average_precision_on_two_sorted_samples_u64_f64", u64, f64, f64, u64, f64, f64, u64, f64, f64, m);
1102
1103    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_bool_f32, "roc_auc_on_two_sorted_samples_bool_f32", bool, f32, f32, bool, f32, f32, bool, f32, f32, m);
1104    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i8_f32, "roc_auc_on_two_sorted_samples_i8_f32", i8, f32, f32, i8, f32, f32, i8, f32, f32, m);
1105    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i16_f32, "roc_auc_on_two_sorted_samples_i16_f32", i16, f32, f32, i16, f32, f32, i16, f32, f32, m);
1106    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i32_f32, "roc_auc_on_two_sorted_samples_i32_f32", i32, f32, f32, i32, f32, f32, i32, f32, f32, m);
1107    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i64_f32, "roc_auc_on_two_sorted_samples_i64_f32", i64, f32, f32, i64, f32, f32, i64, f32, f32, m);
1108    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u8_f32, "roc_auc_on_two_sorted_samples_u8_f32", u8, f32, f32, u8, f32, f32, u8, f32, f32, m);
1109    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u16_f32, "roc_auc_on_two_sorted_samples_u16_f32", u16, f32, f32, u16, f32, f32, u16, f32, f32, m);
1110    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u32_f32, "roc_auc_on_two_sorted_samples_u32_f32", u32, f32, f32, u32, f32, f32, u32, f32, f32, m);
1111    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u64_f32, "roc_auc_on_two_sorted_samples_u64_f32", u64, f32, f32, u64, f32, f32, u64, f32, f32, m);
1112    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_bool_f64, "roc_auc_on_two_sorted_samples_bool_f64", bool, f64, f64, bool, f64, f64, bool, f64, f64, m);
1113    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i8_f64, "roc_auc_on_two_sorted_samples_i8_f64", i8, f64, f64, i8, f64, f64, i8, f64, f64, m);
1114    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i16_f64, "roc_auc_on_two_sorted_samples_i16_f64", i16, f64, f64, i16, f64, f64, i16, f64, f64, m);
1115    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i32_f64, "roc_auc_on_two_sorted_samples_i32_f64", i32, f64, f64, i16, f64, f64, i16, f64, f64, m);
1116    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_i64_f64, "roc_auc_on_two_sorted_samples_i64_f64", i64, f64, f64, i64, f64, f64, i64, f64, f64, m);
1117    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u8_f64, "roc_auc_on_two_sorted_samples_u8_f64", u8, f64, f64, u8, f64, f64, u8, f64, f64, m);
1118    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u16_f64, "roc_auc_on_two_sorted_samples_u16_f64", u16, f64, f64, u16, f64, f64, u16, f64, f64, m);
1119    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u32_f64, "roc_auc_on_two_sorted_samples_u32_f64", u32, f64, f64, u32, f64, f64, u32, f64, f64, m);
1120    roc_auc_on_two_sorted_samples_py!(roc_auc_on_two_sorted_samples_u64_f64, "roc_auc_on_two_sorted_samples_u64_f64", u64, f64, f64, u64, f64, f64, u64, f64, f64, m);
1121
1122    m.add_function(wrap_pyfunction!(loo_cossim_py, m)?).unwrap();
1123    m.add_function(wrap_pyfunction!(loo_cossim_many_py_f64, m)?).unwrap();
1124    m.add_function(wrap_pyfunction!(loo_cossim_many_py_f32, m)?).unwrap();
1125    m.add_class::<PyOrder>().unwrap();
1126    return Ok(());
1127}
1128
1129
1130#[cfg(test)]
1131mod tests {
1132    use super::*;
1133
1134    #[test]
1135    fn test_average_precision_on_sorted() {
1136        let labels: [u8; 4] = [1, 0, 1, 0];
1137        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1138        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1139        let actual: f64 = score_sorted_sample(AveragePrecision::new(), &predictions, &labels, &weights, Order::DESCENDING);
1140        assert_eq!(actual, 0.8333333333333333);
1141    }
1142
1143    #[test]
1144    fn test_average_precision_on_sorted_double() {
1145        let labels: [u8; 8] = [1, 1, 0, 0, 1, 1, 0, 0];
1146        let predictions: [f64; 8] = [0.8, 0.8, 0.4, 0.4, 0.35, 0.35, 0.1, 0.1];
1147        let weights: [f64; 8] = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
1148        let actual: f64 = score_sorted_sample(AveragePrecision::new(), &predictions, &labels, &weights, Order::DESCENDING);
1149        assert_eq!(actual, 0.8333333333333333);
1150    }
1151
1152    #[test]
1153    fn test_average_precision_unsorted() {
1154        let labels: [u8; 4] = [0, 0, 1, 1];
1155        let predictions: [f64; 4] = [0.1, 0.4, 0.35, 0.8];
1156        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1157        let actual: f64 = average_precision(&predictions, &labels, Some(&weights), None);
1158        assert_eq!(actual, 0.8333333333333333);
1159    }
1160
1161    #[test]
1162    fn test_average_precision_sorted() {
1163        let labels: [u8; 4] = [1, 0, 1, 0];
1164        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1165        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1166        let actual: f64 = average_precision(&predictions, &labels, Some(&weights), Some(Order::DESCENDING));
1167        assert_eq!(actual, 0.8333333333333333);
1168    }
1169
1170    #[test]
1171    fn test_average_precision_sorted_pair() {
1172        let labels: [u8; 4] = [1, 0, 1, 0];
1173        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1174        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1175        let actual: f64 = score_two_sorted_samples(
1176            AveragePrecision::new(),
1177            predictions.iter().cloned(),
1178            labels.iter().cloned(),
1179            weights.iter().cloned(),
1180            predictions.iter().cloned(),
1181            labels.iter().cloned(),
1182            weights.iter().cloned()
1183        );
1184        assert_eq!(actual, 0.8333333333333333);
1185    }
1186
1187    #[test]
1188    fn test_roc_auc() {
1189        let labels: [u8; 4] = [1, 0, 1, 0];
1190        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1191        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1192        let actual: f64 = roc_auc(&predictions, &labels, Some(&weights), Some(Order::DESCENDING), None);
1193        assert_eq!(actual, 0.75);
1194    }
1195
1196    #[test]
1197    fn test_roc_auc_double() {
1198        let labels: [u8; 8] = [1, 0, 1, 0, 1, 0, 1, 0];
1199        let predictions: [f64; 8] = [0.8, 0.4, 0.35, 0.1, 0.8, 0.4, 0.35, 0.1];
1200        let actual: f64 = roc_auc(&predictions, &labels, None::<&[f64; 8]>, None, None);
1201        assert_eq!(actual, 0.75);
1202    }
1203
1204    #[test]
1205    fn test_roc_sorted_pair() {
1206        let labels: [u8; 4] = [1, 0, 1, 0];
1207        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1208        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1209        let actual: f64 = score_two_sorted_samples(
1210            RocAuc::new(),
1211            predictions.iter().cloned(),
1212            labels.iter().cloned(),
1213            weights.iter().cloned(),
1214            predictions.iter().cloned(),
1215            labels.iter().cloned(),
1216            weights.iter().cloned()
1217        );
1218        assert_eq!(actual, 0.75);
1219    }
1220
1221    #[test]
1222    fn test_roc_auc_max_fpr() {
1223        let labels: [u8; 4] = [1, 0, 1, 0];
1224        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1225        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1226        let actual: f64 = roc_auc(&predictions, &labels, Some(&weights), Some(Order::DESCENDING), Some(0.25));
1227        assert_eq!(actual, 0.7142857142857143);
1228    }
1229
1230    #[test]
1231    fn test_roc_auc_max_fpr_double() {
1232        let labels: [u8; 8] = [1, 0, 1, 0, 1, 0, 1, 0];
1233        let predictions: [f64; 8] = [0.8, 0.4, 0.35, 0.1, 0.8, 0.4, 0.35, 0.1];
1234        let actual: f64 = roc_auc(&predictions, &labels, None::<&[f64; 8]>, None, Some(0.25));
1235        assert_eq!(actual, 0.7142857142857143);
1236    }
1237
1238    #[test]
1239    fn test_roc_auc_max_fpr_sorted_pair() {
1240        let labels: [u8; 4] = [1, 0, 1, 0];
1241        let predictions: [f64; 4] = [0.8, 0.4, 0.35, 0.1];
1242        let weights: [f64; 4] = [1.0, 1.0, 1.0, 1.0];
1243        let actual: f64 = score_two_sorted_samples(
1244            RocAucWithMaxFPR::new(0.25),
1245            predictions.iter().cloned(),
1246            labels.iter().cloned(),
1247            weights.iter().cloned(),
1248            predictions.iter().cloned(),
1249            labels.iter().cloned(),
1250            weights.iter().cloned()
1251        );
1252        assert_eq!(actual, 0.7142857142857143);
1253    }
1254}
_scors/lib.rs

_scors/
lib.rs
