Skip to main content

CELL_LIST_WGSL

Constant CELL_LIST_WGSL 

Source
pub const CELL_LIST_WGSL: &str = r#"
pub(super) struct CellParams {
    nx:      u32,
    ny:      u32,
    nz:      u32,
    cell_size: f32,
    n_atoms: u32,
    box_x:   f32,
    box_y:   f32,
    box_z:   f32,
}

@group(0) @binding(0) var<storage, read>       positions:    array<vec4<f32>>;
@group(0) @binding(1) var<storage, read_write> cell_indices: array<u32>;
@group(0) @binding(2) var<uniform>             params:       CellParams;

@compute @workgroup_size(64)
pub(super) fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let atom_id = id.x;
    if (atom_id >= params.n_atoms) { return; }

    let pos = positions[atom_id].xyz;

    let ix = clamp(u32(pos.x / params.cell_size), 0u, params.nx - 1u);
    let iy = clamp(u32(pos.y / params.cell_size), 0u, params.ny - 1u);
    let iz = clamp(u32(pos.z / params.cell_size), 0u, params.nz - 1u);

    cell_indices[atom_id] = iz * params.ny * params.nx
                          + iy * params.nx
                          + ix;
}
"#;
Expand description

WGSL compute shader for parallel cell-list construction.