henad-models 0.3.0

Example models for Henad, a parallel agent-based modelling engine.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73


#import henad::rng::{pcg_hash, random_float}
#import henad::space::{cell_index, offset_cell, TORUS}
struct Params {
    width: u32,
    height: u32,
    infection_rate: f32,
    recovery_rate: f32,
}

// --8<-- [start:bindings]
@group(0) @binding(0) var<storage, read> state_in: array<u32>;
@group(0) @binding(1) var<storage, read_write> state_out: array<u32>;
@group(0) @binding(2) var<storage, read> rng_in: array<u32>;
@group(0) @binding(3) var<storage, read_write> rng_out: array<u32>;
@group(0) @binding(4) var<uniform> params: Params;
// --8<-- [end:bindings]

const S: u32 = 0u;
const I: u32 = 1u;
const R: u32 = 2u;

// Each cell keeps its `pcg_hash` state in the ping-ponged `rng` buffer, and advances it one round per step.
// A seed derived from a `tick` uniform would repeat across a batch. A uniform written between the passes of one
// encoder takes effect only when the encoder submits, and every pass of the batch would read the same tick.

@compute
@workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let width = params.width;
    let height = params.height;
    if (global_id.x >= width || global_id.y >= height) {
        return;
    }

    let x = global_id.x;
    let y = global_id.y;
    let idx = cell_index(x, y, width);

    let left = u32(offset_cell(x, y, -1, 0, width, height, TORUS).x);
    let right = u32(offset_cell(x, y, 1, 0, width, height, TORUS).x);
    let up = u32(offset_cell(x, y, 0, -1, width, height, TORUS).y);
    let down = u32(offset_cell(x, y, 0, 1, width, height, TORUS).y);

    var infected_neighbors: u32 = 0u;
    infected_neighbors += u32(state_in[cell_index(left, up, width)] == I);
    infected_neighbors += u32(state_in[cell_index(x, up, width)] == I);
    infected_neighbors += u32(state_in[cell_index(right, up, width)] == I);
    infected_neighbors += u32(state_in[cell_index(left, y, width)] == I);
    infected_neighbors += u32(state_in[cell_index(right, y, width)] == I);
    infected_neighbors += u32(state_in[cell_index(left, down, width)] == I);
    infected_neighbors += u32(state_in[cell_index(x, down, width)] == I);
    infected_neighbors += u32(state_in[cell_index(right, down, width)] == I);

    let rng_next = pcg_hash(rng_in[idx]);
    rng_out[idx] = rng_next;

    let cell = state_in[idx];
    var next_cell = cell;
    if (cell == S && infected_neighbors > 0u) {
        let prob_safe = pow(1.0 - params.infection_rate, f32(infected_neighbors));
        // `>=`, matching the CPU model. A zero draw would otherwise escape a rate of 1.
        if (random_float(rng_next, 1.0) >= prob_safe) {
            next_cell = I;
        }
    } else if (cell == I) {
        if (random_float(rng_next, 1.0) < params.recovery_rate) {
            next_cell = R;
        }
    }

    state_out[idx] = next_cell;
}