r2rs-nmath 0.1.1

Statistics programming for Rust based on R's nmath package
Documentation
// Translation of nmath's cwilcox
//
// Copyright (C) 2020 neek-sss
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use std::{cell::RefCell, thread_local};

use fnv::FnvHashMap;

thread_local! {
    static W: RefCell<FnvHashMap<usize, FnvHashMap<usize, FnvHashMap<usize, f64>>>> = RefCell::new(FnvHashMap::default());
}

pub fn reset_wilcox_statics() {
    W.with(|W_| *W_.borrow_mut() = FnvHashMap::default());
}

/// This counts the number of choices with statistic = k
pub fn cwilcox(mut k: i32, m: usize, n: usize) -> f64 {
    let mut c = 0; /* hence  k <= (u / 2).floor() */
    let mut u = 0; /* hence  i <= j */
    let mut i = 0;
    let mut j = 0;
    u = m * n;
    if k < 0 || k > u as i32 {
        return 0.0;
    }
    c = u / 2;
    if k > c as i32 {
        k = u as i32 - k; /* hence  k <= floor(u / 2) */
    }
    if m < n {
        i = m;
        j = n;
    } else {
        i = n;
        j = m;
    } /* hence  i <= j */

    if j == 0 {
        /* and hence i == 0 */
        return if k == 0 { 1.0 } else { 0.0 };
    }

    /* We can simplify things if k is small.  Consider the Mann-Whitney
       definition, and sort y.  Then if the statistic is k, no more
       than k of the y's can be <= any x[i], and since they are sorted
       these can only be in the first k.  So the count is the same as
       if there were just k y's.
    */
    if j > 0 && k < j as i32 {
        return cwilcox(k, i, k as usize);
    }

    let mut w_current = W.with(|W_| {
        let mut w = W_.borrow_mut();
        *w.entry(i)
            .or_insert_with(FnvHashMap::default)
            .entry(j)
            .or_insert_with(FnvHashMap::default)
            .entry(k as usize)
            .or_insert(-1.0)
    });

    w_current = if w_current < 0.0 {
        let w_temp = if j == 0 {
            /* and hence i == 0 */
            if k == 0 {
                1.0
            } else {
                0.0
            }
        } else {
            cwilcox(k - j as i32, i - 1, j) + cwilcox(k, i, j - 1)
        };

        W.with(|W_| {
            let mut w = W_.borrow_mut();
            *w.entry(i)
                .or_insert_with(FnvHashMap::default)
                .entry(j)
                .or_insert_with(FnvHashMap::default)
                .entry(k as usize)
                .or_insert(-1.0) = w_temp;
        });

        w_temp
    } else {
        w_current
    };

    return w_current;
}