Epidemiological simulation in V

sir.v

// v -cc gcc -cflags "-O3 -ffast-math -flto" sir.v
import math
import time

[inline]
fn clz64(v u64) u32 {
    return C.__builtin_clzll(v)
}

[inline]
fn xorshift128plus(rng mut []u64) u64 {
    mut x := rng[0]
    y := rng[1]
    rng[0] = y
    x ^= x << 23
    rng[1] = u64(x) ^ u64(y) ^ u64(x >> u64(17)) ^ u64(u64(y) >> u64(26))
    z := rng[1] + y
    return z
}

[inline]
fn randunif(rng mut []u64) f64 {
    mut significand := xorshift128plus(mut rng)
    mut exponent := int(-64)
    for significand == 0 {
        exponent -= 64
        if exponent < -1074 {
            return f64(0.0)
        }
        significand = xorshift128plus(mut rng)
    }
    shift := clz64(significand)
    if shift != 0 {
        exponent -= int(shift)
        significand <<= u64(shift)
        significand |= u64(significand >> u64(64 - int(shift)))
    }
    significand |= u64(1)
    return C.ldexp(f64(significand),exponent)
}

[inline]
fn randbn(n i64, p f64, rng mut []u64 ) i64 {
    log_q := math.log(1.0-p)
    mut x := i64(0)
    mut sum := f64(0.0)
    for sum >= log_q {
        sum += math.log(randunif(mut rng))/f64(n-x)
        if sum < log_q {
            break
        }
        x += 1
    }
    return x
}

[inline]
fn sir(u []i64, du mut []i64, parms []f64, rng mut []u64) {
    s := u[0]
    i := u[1]
    r := u[2]
    y := u[3]
    beta := parms[0]
    gamma := parms[1]
    iota := parms[2]
    n := parms[3]
    dt := parms[4]
    lambd := beta*(f64(i)+iota)/n
    ifrac := 1.0-math.exp(-lambd*dt)
    rfrac := 1.0-math.exp(-gamma*dt)
    infection := randbn(s, ifrac, mut rng)
    recovery := randbn(i, rfrac, mut rng)
    du[0] = s - infection
    du[1] = i + infection - recovery
    du[2] = r + recovery
    du[3] = y + infection
}

[inline]
fn sum(y []i64) i64 {
    mut sm := i64(0)
    for x in y {
        sm += x
    }
    return sm
}

fn simulate() f64 {
    parms := [f64(2.62110617498984), 0.5384615384615384, 0.5, 403.0, 0.1]
    tf := 540
    nsims := 1000
    mut yvec := [i64(0)].repeat(nsims)
    mut rng := [u64(123),456]
    mut i := 0
    mut j := 1
    mut u := [i64(0),0,0,0]
    mut du := [i64(0),0,0,0]
    i = 0
    for i < nsims {
        u = [i64(60),1,342,0]
        du = [i64(0),0,0,0]
        j = 1
        for j <= tf {
            sir(u, mut du, parms, mut rng)
            u = du
            j++
        }
        yvec[i] = u[3]
        i++
    }
    numerator := sum(yvec)
    result := f64(numerator)/f64(nsims)
    return result
}

fn main(){
    t0 := time.ticks()
    m := simulate()
    t1 := time.ticks()
    tdiff := f32(t1-t0)/1000.0
    println('$tdiff')
    println('$m')
}