Epidemiological simulation in Rust

sir.rs

#!/usr/bin/env run-cargo-script
//!
//! ```cargo
//! [dependencies]
//! time = "0.1.25"
//! float_extras = "0.1.1"
//! xorshift = "0.1"
//! ```

#![allow(non_snake_case)]

use std::f64;
use std::u64;
use std::i32;

extern crate xorshift;
use xorshift::{Rng, SeedableRng, Xorshift128};
extern crate float_extras;

#[inline]
fn random_real(x: u64) -> f64 {
    let mut exponent: i32 = -64i32;
    let mut significand: u64;
    let shift: u32;
    loop  {
        significand = x;
        if !(significand == 0) { break ; }
        exponent -= 64i32;
        if exponent < -1074i32 { return 0.0 }
    }
    shift = significand.leading_zeros();
    if shift != 0 {
        exponent -= shift as i32;
        significand <<= shift;
        significand |= x >> (64 - shift);
    }
    significand |= 1;
    float_extras::f64::ldexp(significand as f64, exponent as isize)
}

#[inline]
fn randunif(rng: &mut Xorshift128) -> f64 {
    let z: u64 = rng.next_u64();
    let dz: f64 = random_real(z);
    dz
}

#[inline]
fn randbn(n: i64, p: f64, mut rng: &mut Xorshift128) -> i64 {
    let log_q: f64 = f64::ln(1.0f64 - p);
    let mut x: i64 = 0;
    let mut sum: f64 = 0.0;
    loop  {
        sum += f64::ln(randunif(&mut rng)) / ((n - x) as f64);
        if sum < log_q { break ; }
        x += 1;
    }
    x
}

#[inline]
fn sir(u: &mut [i64; 4], du: &mut [i64; 4],
    parms: &[f64; 5], mut rng: &mut Xorshift128) {
    let S: i64 = u[0];
    let I: i64 = u[1];
    let R: i64 = u[2];
    let Y: i64 = u[3];
    let beta: f64 = parms[0];
    let gamma: f64 = parms[1];
    let iota: f64 = parms[2];
    let N: f64 = parms[3];
    let dt: f64 = parms[4]; 
    let lambd: f64 = beta * (I as f64 + iota) / N;
    let ifrac: f64 = 1.0f64 - f64::exp(-lambd * dt);
    let rfrac: f64 = 1.0f64 - f64::exp(-gamma * dt);
    let infection: i64 = randbn(S, ifrac, &mut rng);
    let recovery: i64 = randbn(I, rfrac, &mut rng);
    du[0] = S - infection;
    du[1] = I + infection - recovery;
    du[2] = R + recovery;
    du[3] = Y + infection;
}

fn simulate() -> f64 {
    let parms: [f64; 5] = [2.62110617498984f64, 0.5384615384615384f64,
                        0.5f64, 403.0f64, 0.1f64];
    let tf: i64 = 540;
    let nsims: i64 = 1000;
    let mut numerator: i64 = 0;
    let states = [123, 456];
    let mut rng: Xorshift128 = SeedableRng::from_seed(&states[..]);
    let mut i: i64 = 0;
    while i < nsims {
        let mut u: [i64; 4] = [60, 1, 342, 0];
        let mut du: [i64; 4] = [0, 0, 0, 0];
        let mut j: i64 = 1;
        while j <= tf {
            sir(&mut u, &mut du, &parms, &mut rng);
            u = du;
            j += 1
        }
        numerator += u[3];
        i += 1
    }
    (numerator as f64) / (nsims as f64)
}

pub fn main() {
    let tic = time::precise_time_ns();
    let m = simulate();
    let toc = time::precise_time_ns();
    let interval = toc-tic;
    println!("{}\n{}\n", (interval as f64)/1000000000.0, m);
}