Neural mass network simulations in Futhark

mpr-append.fut

-- this variant uses functional approach, appending new states to buffer
-- instead of in-place update

let pi = 3.141592653589793f32

-- some type abbreviations
type mpr_pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type mpr_node = (f32, f32)
type mpr_net [n] = [n] mpr_node
-- this is tranposed from mpr-pdq to avoid tranposes in history update
type mpr_hist [t] [n] = [t] mpr_net [n]
type connectome [n] = {weights: [n][n]f32, idelays: [n][n]i64}

-- do one time step w/ Euler
let mpr_step [t] [n] (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): []mpr_net[n] =

    -- define individual derivatives as in mpr pdq
    let dr r V = 1/p.tau * (p.Delta / (pi * p.tau) + 2 * V * r)
    let dV r V r_c = 1/p.tau * ( V**2 - pi**2 * p.tau**2 * r**2 + p.eta + p.J * p.tau * r + p.I + r_c)
    let dfun (r, V, c) = (dr r V, dV r V c)

    -- unpack current state for clarity
    let (r, V) = last buf |> unzip

    -- connectivity eval
    let r_c_i i w d = map2 (\wj dj -> wj * buf[t - dj - 1, i].0) w d |> reduce (+) 0f32 |> (*p.G)
    let r_c = map3 r_c_i (iota n) conn.weights conn.idelays
    
    -- Euler step
    let erV = map3 (\r V c -> (dr r V, dV r V c)) r V r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))

    -- now for the Heun step
    let (er, eV) = unzip erV
    let hrV = map3 (\r V c -> (dr r V, dV r V c)) er eV r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))    

    -- return updated buffer
    in buf ++ [hrV]

let run_mpr [t] [n]
        (horizon: i64) (nt: i64) (tt: i64) (dt: f32)
        (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): mpr_hist[tt][n] =
    (loop (buf) for now < nt do mpr_step dt buf conn p) :> mpr_hist[tt][n]

let mpr_pars_with_G (p: mpr_pars) (new_G: f32): mpr_pars = 
    let new_p = copy p
    in new_p with G = new_G

let loss [t] [n] (x:mpr_hist[t][n]): f32 =
    let r = map unzip x[t-10:] |> unzip |> (.0)
    let sum = map (reduce (+) 0f32) r |> reduce (+) 0f32
    in
    sum

let sweep [t] [n]
        (ng: i64) (horizon: i64) (nt: i64) (tt: i64) (dt: f32)
        (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): [ng]f32 =
    let Gs = tabulate ng (\i -> 0.0 + (f32.i64 i) * 0.1)
    let do_one G = (run_mpr horizon nt (nt + horizon + 1) dt buf conn (mpr_pars_with_G p G)) |> loss 
    in map (\g -> vjp do_one g 1f32) Gs

-- ==
-- input { 1i64 256i64 100i64 400i64 }
-- input { 2i64 256i64 100i64 400i64 }
-- input { 4i64 256i64 100i64 400i64 }
-- input { 8i64 256i64 100i64 400i64 }
let main (ng: i64) (nh: i64) (nt: i64) (nn: i64) =
    let dt = 0.01f32
    let buf = tabulate_2d (nh + 1) nn (\i j -> (0.1f32, -2.0f32))
    let conn = {weights=tabulate_2d nn nn (\i j -> 0.1f32),
                idelays=tabulate_2d nn nn (\i j -> (i * j) % nh)}
    let p = {G=0.1f32, I=0.0f32, Delta=0.7f32, eta=(-4.6f32), tau=1.0f32, J=14.5f32}
    in sweep ng nh nt (nt + nh + 1) dt buf conn p
    -- in run_mpr nh nt (nt + nh + 1) dt buf conn p

mpr-inplace.fut

let pi = 3.141592653589793f32

-- some type abbreviations
type mpr_pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type mpr_node = (f32, f32)
type mpr_net [n] = [n] mpr_node
-- this is tranposed from mpr-pdq to avoid tranposes in history update
type mpr_hist [t] [n] = [t] mpr_net [n]
type connectome [n] = {weights: [n][n]f32, idelays: [n][n]i64}

-- do one time step w/ Euler
let mpr_step [t] [n] (now: i64) (dt: f32) (buf: *mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): *mpr_hist[t][n] =

    -- define individual derivatives as in mpr pdq
    let dr r V = 1/p.tau * ( p.Delta / (pi * p.tau) + 2 * V * r)
    let dV r V r_c = 1/p.tau * ( V**2 - pi**2 * p.tau**2 * r**2 + p.eta + p.J * p.tau * r + p.I + r_c)
    let dfun (r, V, c) = (dr r V, dV r V c)

    -- unpack current state for clarity
    let (r, V) = last buf |> unzip

    -- connectivity eval
    let r_c_i i w d = map2 (\wj dj -> wj * buf[now - dj, i].0) w d |> reduce (+) 0f32 |> (*p.G)
    let r_c = map3 r_c_i (iota n) conn.weights conn.idelays

    -- Euler step
    let erV = map3 (\r V c -> (dr r V, dV r V c)) r V r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))

    -- now for the Heun step
    let (er, eV) = unzip erV
    let hrV = map3 (\r V c -> (dr r V, dV r V c)) er eV r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))    

    -- return updated buffer
    in buf with [now + 1] = (copy hrV)

let run_mpr [t] [n] (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): mpr_hist[t][n] =
    loop buf = copy buf for now < (t - horizon - 1) do mpr_step (now + horizon) dt buf conn p

let mpr_pars_with_G (p: mpr_pars) (new_G: f32): mpr_pars = 
    let new_p = copy p
    in new_p with G = new_G

let sweep [t] [n] (ng: i64) (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): [ng]mpr_hist[t][n] =
    let Gs = tabulate ng (\i -> 0.0 + (f32.i64 i) * 0.1)
    let do_one G = run_mpr horizon dt buf conn (mpr_pars_with_G p G)
    in map do_one Gs

-- ==
-- input { 1i64 256i64 100i64 400i64 }
-- input { 2i64 256i64 100i64 400i64 }
-- input { 4i64 256i64 100i64 400i64 }
-- input { 8i64 256i64 100i64 400i64 }
let main (ng: i64) (nh: i64) (nt: i64) (nn: i64) =
    let dt = 0.01f32
    let buf = tabulate_2d (nt + nh) nn (\i j -> (0.1f32, -2.0f32))
    let conn = {weights=tabulate_2d nn nn (\i j -> 0.1f32),
                idelays=tabulate_2d nn nn (\i j -> (i * j) % nh)}
    let p = {G=0.1f32, I=0.0f32, Delta=0.7f32, eta=(-4.6f32), tau=1.0f32, J=14.5f32}
    in sweep ng nh dt buf conn p
    -- in run_mpr nh dt buf conn p

mpr.fut

-- revised in-place version, best for cpu/multicore

let pi = 3.141592653589793f32

-- some type abbreviations
type mpr_pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type mpr_node = (f32, f32)
type mpr_net [n] = [n] mpr_node
-- this is tranposed from mpr-pdq to avoid tranposes in history update
type mpr_hist [t] [n] = [t] mpr_net [n]
type connectome [n] = {weights: [n][n]f32, idelays: [n][n]i64}

-- do one time step w/ Euler
let mpr_step [t] [n] (now: i64) (dt: f32) (buf: *mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): *mpr_hist[t][n] =

    -- define individual derivatives as in mpr pdq
    let dr r V = 1/p.tau * ( p.Delta / (pi * p.tau) + 2 * V * r)
    let dV r V r_c = 1/p.tau * ( V**2 - pi**2 * p.tau**2 * r**2 + p.eta + p.J * p.tau * r + p.I + r_c)
    let dfun (r, V, c) = (dr r V, dV r V c)

    -- unpack current state for clarity
    let (r, V) = last buf |> unzip

    -- connectivity eval
    let r_c_i i w d = map2 (\wj dj -> wj * buf[now - dj, i].0) w d |> reduce (+) 0f32 |> (*p.G)
    let r_c = map3 r_c_i (iota n) conn.weights conn.idelays

    -- Euler step
    let erV = map3 (\r V c -> (dr r V, dV r V c)) r V r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))

    -- now for the Heun step
    let (er, eV) = unzip erV
    let hrV = map3 (\r V c -> (dr r V, dV r V c)) er eV r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))    

    -- return updated buffer
    in buf with [now + 1] = copy hrV

let run_mpr [t] [n] (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): mpr_hist[t][n] =
    loop buf = copy buf
        for now < (t - horizon - 1) do mpr_step (now + horizon) dt buf conn p

let mpr_pars_with_G (p: mpr_pars) (new_G: f32): mpr_pars = 
    let new_p = copy p
    in new_p with G = new_G

let sweep [t] [n] (ng: i64) (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): [ng]mpr_hist[t][n] =
    let Gs = tabulate ng (\i -> 0.0 + (f32.i64 i) * 0.1)
    let do_one G = run_mpr horizon dt buf conn (mpr_pars_with_G p G)
    in map do_one Gs

-- input { 16384i64 256i64 1i64 164i64 }

-- input { 2i64 512i64 1i64 76i64 }
-- input { 4i64 512i64 1i64 76i64 }
-- input { 8i64 512i64 1i64 76i64 }
-- input { 16i64 512i64 1i64 76i64 }
-- input { 32i64 512i64 1i64 76i64 }
-- input { 64i64 512i64 1i64 76i64 }
-- input { 128i64 512i64 1i64 76i64 }

-- ==
-- input { 1i64 512i64 10000i64 76i64 }
-- input { 2i64 512i64 10000i64 76i64 }
-- input { 4i64 512i64 10000i64 76i64 }
-- input { 8i64 512i64 10000i64 76i64 }
let main (ng: i64) (nh: i64) (nt: i64) (nn: i64) =
    let dt = 0.01f32
    let buf = tabulate_2d (nt + nh) nn (\i j -> (0.1f32, -2.0f32))
    let conn = {weights=tabulate_2d nn nn (\i j -> 0.1f32),
                idelays=tabulate_2d nn nn (\i j -> ((i * j) % nh))
                }
    let p = {G=0.1f32, I=0.0f32, Delta=0.7f32, eta=(-4.6f32), tau=1.0f32, J=14.5f32}
    in sweep ng nh dt buf conn p
    -- in run_mpr nh dt buf conn p

-- on CPU (4 core w/ boost 8 MB cache)
-- g=128   6618 us   51 us/   10.4 MB

-- CPU saturates around 30k steps/s here, 
-- nb_pdq is 20 ms for 100ms sim @ dt=0.1 i.e. 1000 steps, i.e. 20 us/ (but nn==76)

-- so about 20% perf hit on CPU for sweeps, 50% on GPU, with some caveats:
-- -> pdq uses a better layout, maybe worth implementing since the write isn't costly?
-- -> pdq bench running on single core, with a boost not present for single core
-- -> single-threaded C Futhark is significantly slower
-- yet no porting required between C, multicore & GPU

-- https://futhark-lang.org/student-projects/duc-msc-thesis.pdf
-- details the multicore backend, task scheduler, avoiding false sharing etc etc etc
-- in summary, it's very well thought out, probably producing nicely SIMD code?
-- TODO check explicitly SIMD code in the multicore backend?

-- hm, without the outer map, and many time steps, we've 737942us/100000 steps,
-- which is *7* us / step; w/ 1k steps, 9 us per step. (CPU boost to 3.8 GHz)
-- so we can call it 2x Numba perf for single thread?
-- removing noise and using f32 in number is 15% improvement

-- now with more time steps on multicore, g=4, 197ms/10k steps, ~20us/s
-- which matches numba speed for single sim.

-- on GPU mpr1 t>1 performs worse even after tuning: 2.5ms 1 step, 32 ms 10 steps??
-- maybe lifting the time loop is good there?
-- mpr4 does better for multiple steps, g=1024, 18 ms 10 steps (568k step/s)
-- , 148 ms for 100 steps -> 700k steps/s (saturates)

-- this is factor ~2x slower than (my) hand written CUDA (w/o shuffles)




-- on GPU this scales not so bad once tuning is done
-- and it scales better than the lifted time loop
-- g=1        84us  84  us/
-- g=1024   2739us   2.7us/
-- g=8192  21586us   2.6us/
-- g=16384 43360us   2.6us/

-- this is saturating at ~400k step per second, less than full CUDA by 5x or so
-- but still 10x faster than 6 core CPU

-- why using 4-5 GB ? 16384*256*164*2 svar*4b/f32-> 5.4 GB

-- manually lifting of copy to buf2 produces same results (smart compiler)

mpr4.fut

-- pulls the time loop outside the parallel map, best for GPU
let pi = 3.141592653589793f32

-- some type abbreviations
type pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type node = (f32, f32)
type net [n] = [n] node
-- this is tranposed from mpr-pdq to avoid tranposes in history update
type hist [t] [n] = [t] net [n]
type conn [n] = {weights: [n][n]f32, idelays: [n][n]i64}

-- compute next time step
let step [t] [n] (dt: f32) (buf: hist[t][n]) (c: conn[n]) (p: pars): net[n] =

    -- define individual derivatives as in mpr pdq
    let dr r V = 1/p.tau * (p.Delta / (pi * p.tau) + 2 * V * r)
    let dV r V r_c = 1/p.tau * ( V**2 - pi**2 * p.tau**2 * r**2 + p.eta + p.J * p.tau * r + p.I + r_c)
    let dfun (r, V, c) = (dr r V, dV r V c)

    -- unpack current state for clarity
    let (r, V) = last buf |> unzip

    -- connectome eval
    let r_c_i i w d = map2 (\wj dj -> wj * buf[t - dj - 1, i].0) w d |> reduce (+) 0f32 |> (*p.G)
    let r_c = map3 r_c_i (iota n) c.weights c.idelays

    -- Euler step
    let erV = map3 (\r V c -> (dr r V, dV r V c)) r V r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))

    -- Heun step
    let (er, eV) = unzip erV
    let hrV = map3 (\r V c -> (dr r V, dV r V c)) er eV r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))    

    -- return next state
    in hrV

let sweep [g] [t] [n] (dt:f32) 
        (buf: [g]hist[t][n]) (c: conn[n]) (p: [g]pars): [g]net[n] =
    map2 (\b p -> step dt b c p) buf p

let run [g] [t] [n] (T:i64) (dt:f32) 
        (buf: hist[t][n]) (c: conn[n]) (p: [g]pars): [g]hist[t][n] =
    -- should be linear (?)
    let buf = tabulate g (\_ -> copy buf) -- alloc full storage
    let buf =
        loop (buf) for i < T do
            let next = sweep dt buf c p
            let buf[:,t-1] = copy next
            in buf
    in buf



-- requires tuning with T=100

-- ==
-- input { 1024i64 512i64 1i64 76i64 }
-- input { 1024i64 512i64 10i64 76i64 }
-- input { 1024i64 512i64 100i64 76i64 }
-- input { 1024i64 512i64 200i64 76i64 }
let main (ng: i64) (nh: i64) (T: i64) (nn: i64) =
    let dt = 0.01f32
    let buf = tabulate_2d (nh + 1) nn (\i j -> (0.1f32, -2.0f32))
    let conn = {weights=tabulate_2d nn nn (\i j -> 0.1f32),
                idelays=tabulate_2d nn nn (\i j -> (i * j) % nh)}
    let p = {G=0.1f32, I=0.0f32, Delta=0.7f32, eta=(-4.6f32), tau=1.0f32, J=14.5f32}
    let p = tabulate ng (\_ -> p)
    in run T dt buf conn p

-- works fine on gpu w/ tuning

mpr6.fut

-- see notes below.  this version works well on CPU & GPU, default to seq=true

let pi = 3.141592653589793f32

-- lifted the body of n < N loop for use with both map & loop
let mpr_node [N][T] (i: i64) (n: i64) (dt: f32) (nstep: i64) (i0: i64)
        (r: [N][T]f32) (V: [N][T]f32) (weights: [N][N]f32) (idelays: [N][N]i64)
        (G: f32) (I: f32) (Delta: f32) (eta: f32) (tau: f32) (J: f32):
        -- returns next r, V
        (f32, f32) =

    let dr r V = 1/tau * ( Delta / (pi * tau) + 2 * V * r)
    let dV r V r_c = 1/tau * ( V**2 - pi**2 * tau**2 * r**2 + eta + J * tau * r + I + r_c)
    let r_bound r = if r >= 0f32 then r else 0f32

    -- coupling
    let r_c = iota N |> map (\m -> weights[n,m] * r[m,i - idelays[n,m] - 1]) |> reduce (+) 0f32  
    let r_c = r_c * G

    -- TODO precomputed additive noise

    -- Heun integration step
    let dr_0 = dr r[n,i-1] V[n,i-1]
    let dV_0 = dV r[n,i-1] V[n,i-1] r_c

    let r_int = r[n,i-1] + dt*dr_0
    let V_int = V[n,i-1] + dt*dV_0
    let r_int = r_bound r_int

    let r_next = r[n,i-1] + dt*(dr_0 + (dr r_int V_int))/2f32
    let V_next = V[n,i-1] + dt*(dV_0 + (dV r_int V_int r_c))/2f32
    let r_next = r_bound r_next

    -- loops return next (r, V)
    in (r_next, V_next)   

-- this version forces sequential N loop to keep GPU perf high
let mpr_integrate_seq [N] [T] (dt: f32) (nstep: i64) (i0: i64)
        (r: *[N][T]f32) (V: *[N][T]f32) (weights: [N][N]f32) (idelays: [N][N]i64)
        (G: f32) (I: f32) (Delta: f32) (eta: f32) (tau: f32) (J: f32): 
        -- returns r, V updated
        (*[N][T]f32, *[N][T]f32) =
    loop (r, V) for i_ < nstep do
        loop (r, V) for n < N do
            let i = i_ + i0
            let (rn, Vn) = mpr_node i n dt nstep i0 r V weights idelays G I Delta eta tau J
            let r[n,i] = rn
            let V[n,i] = Vn
            in (r, V)

-- this version allows parallel N loop
let mpr_integrate_map [N] [T] (dt: f32) (nstep: i64) (i0: i64)
        (r: *[N][T]f32) (V: *[N][T]f32) (weights: [N][N]f32) (idelays: [N][N]i64)
        (G: f32) (I: f32) (Delta: f32) (eta: f32) (tau: f32) (J: f32): 
        -- returns r, V updated
        (*[N][T]f32, *[N][T]f32) =
    loop (r, V) for i_ < nstep do -- time loop, sequential
        let i = i_ + i0
        let f n = mpr_node i n dt nstep i0 r V weights idelays G I Delta eta tau J
        let (nr, nV) = map f (iota N) |> unzip
        let r[:,i] = nr
        let V[:,i] = nV
        in (r, V)

let sweep [N] [T] (g: i64) (dt: f32) (nstep: i64) (i0: i64)
        (r: [N][T]f32) (V: [N][T]f32) 
        (weights: [N][N]f32) (idelays: [N][N]i64)
        (seq: bool): 
            [g]([N][T]f32, [N][T]f32) =
    let Gs = tabulate g (\i -> 0.0 + (f32.i64 i) * 0.1)
    -- TODO better way to choose a function in a branch?
    let do_one_seq G = mpr_integrate_seq dt nstep i0 (copy r) (copy V) weights idelays G 0.0f32 0.7f32 (-4.6f32) 1.0f32 14.5f32
    let do_one_map G = mpr_integrate_map dt nstep i0 (copy r) (copy V) weights idelays G 0.0f32 0.7f32 (-4.6f32) 1.0f32 14.5f32
    in if seq then map do_one_seq Gs else map do_one_map Gs

-- GPU sizes, tune & bench
-- input { 256i64 256i64 10i64 68i64 }
-- input { 32768i64 256i64 200i64 68i64 }

-- CPU sizes

-- ==
-- input { 1i64 256i64 1000i64 68i64 true }
-- input { 64i64 256i64 1000i64 68i64 true }
-- input { 1i64 256i64 1000i64 68i64 false }
-- input { 64i64 256i64 1000i64 68i64 false }
let main (ng: i64) (nh: i64) (nt: i64) (nn: i64) (seq: bool) =
    let dt = 0.01f32
    let r = tabulate_2d nn (nh + nt) (\i j -> 0.1f32)
    let V = tabulate_2d nn (nh + nt) (\i j -> -2.0f32)
    let weights = tabulate_2d nn nn (\i j -> 0.1f32)
    let idelays = tabulate_2d nn nn (\i j -> ((i * j) % nh))
    in sweep ng dt nt nh r V weights idelays seq

-- with 32k, 12 GB mem used, 180W/230W 436 ms per iter, 7.5M/s
-- with { 32768i64 256i64 200i64 68i64 } hit 190W, 8.6M/s
-- requires tuning with { 256i64 256i64 10i64 68i64 }
-- [n][t] vs [t][n] doesn't matter for GPU but 5% CPU. flat writes could be good.

-- can hit 1.1M/s with { 96i64 256i64 1000i64 68i64 true } on CPU
-- for g=1, Xeon W +40% seq, mba +2x map, epyc gen2 +3x seq

-- so this is done for now, even if we might want to try
-- the transpose for a flat history write

mprad.fut

-- test mpr sim with ad for params

let pi = 3.141592653589793f32

-- some type abbreviations
type mpr_pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type mpr_node = (f32, f32)
type mpr_net [n] = [n] mpr_node
-- this is tranposed from mpr-pdq to avoid tranposes in history update
type mpr_hist [t] [n] = [t] mpr_net [n]
type connectome [n] = {weights: [n][n]f32, idelays: [n][n]i64}

-- do one time step w/ Euler
let mpr_step [t] [n] (now: i64) (dt: f32) (buf: *mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): *mpr_hist[t][n] =

    -- define individual derivatives as in mpr pdq
    let dr r V = 1/p.tau * ( p.Delta / (pi * p.tau) + 2 * V * r)
    let dV r V r_c = 1/p.tau * ( V**2 - pi**2 * p.tau**2 * r**2 + p.eta + p.J * p.tau * r + p.I + r_c)
    let dfun (r, V, c) = (dr r V, dV r V c)

    -- unpack current state for clarity
    let (r, V) = last buf |> unzip

    -- connectivity eval
    let r_c_i i w d = map2 (\wj dj -> wj * buf[now - dj, i].0) w d |> reduce (+) 0f32 |> (*p.G)
    let r_c = map3 r_c_i (iota n) conn.weights conn.idelays

    -- Euler step
    let erV = map3 (\r V c -> (dr r V, dV r V c)) r V r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))

    -- now for the Heun step
    let (er, eV) = unzip erV
    let hrV = map3 (\r V c -> (dr r V, dV r V c)) er eV r_c
           |> map2 (\(r, V) (dr, dV) -> (r + dt * dr, V + dt * dV)) (last buf)
           |> map1 (\(r, V) -> (if r >= 0f32 then r else 0f32, V))    

    -- return updated buffer
    in buf with [now + 1] = copy hrV

let run_mpr [t] [n] (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): mpr_hist[t][n] =
    loop buf = copy buf
        for now < (t - horizon - 1) do mpr_step (now + horizon) dt buf conn p

let mpr_pars_with_G (p: mpr_pars) (new_G: f32): mpr_pars = 
    let new_p = copy p
    in new_p with G = new_G

let loss [t] [n] (x:mpr_hist[t][n]): f32 =
    let r = map unzip x[t-10:] |> unzip |> (.0)
    let sum = map (reduce (+) 0f32) r |> reduce (+) 0f32
    in
    sum

let sweep [t] [n] (ng: i64) (horizon: i64) (dt: f32) (buf: mpr_hist[t][n]) (conn: connectome[n]) (p: mpr_pars): [ng]f32 =
    let Gs = tabulate ng (\i -> 0.0 + (f32.i64 i) * 0.1)
    let do_one G = run_mpr horizon dt buf conn (mpr_pars_with_G p G) |> loss
    in map (\g -> vjp do_one g 1f32) Gs

-- ==
-- input { 1i64 512i64 10000i64 76i64 }
-- input { 2i64 512i64 10000i64 76i64 }
-- input { 4i64 512i64 10000i64 76i64 }
-- input { 8i64 512i64 10000i64 76i64 }
-- input { 16i64 512i64 10000i64 76i64 }
let main (ng: i64) (nh: i64) (nt: i64) (nn: i64) =
    let dt = 0.01f32
    let buf = tabulate_2d (nt + nh) nn (\i j -> (0.1f32, -2.0f32))
    let conn = {weights=tabulate_2d nn nn (\i j -> 0.1f32),
                idelays=tabulate_2d nn nn (\i j -> ((i * j) % nh))
                }
    let p = {G=0.1f32, I=0.0f32, Delta=0.7f32, eta=(-4.6f32), tau=1.0f32, J=14.5f32}
    in sweep ng nh dt buf conn p

mistake here, the Heun step here isn't correct, it's just a second Euler step, but it's a minor modification and doesn't affect performance