qualia_core_db/wgsl_forge/physics/
kinematics.rs

1//! Softened inverse-square N-body step as a certified forge kernel.
2//!
3//! Embeds [`shaders/kinematics.wgsl`](../../../shaders/kinematics.wgsl) via
4//! `include_str!` (single source of truth), grades it against the exact CPU oracle
5//! [`nbody_step_cpu`], and runs it on any wgpu adapter via [`nbody_step_gpu`].
6//!
7//! State is a flat `f32` buffer, 8 scalars per particle:
8//! `[px, py, pz, vx, vy, vz, mass, charge]`. The kernel is **double-buffered**: forces
9//! are read from the input state only and written to a separate output, so the result is
10//! independent of invocation order (no read/write race). The pairwise force on `i` is
11//! `F_i = coupling · Σ_{j≠i} q_i q_j (x_i−x_j) / (|x_i−x_j|² + soft)^{3/2}` (Plummer
12//! softening, no singular skip branch), then a symplectic-Euler update
13//! `v ← v + (F/m)dt`, `x ← x + v·dt`. `coupling` selects the law: `+k` electrostatic
14//! (repulsive like-charges), `−G` gravitational (put mass in the charge slot).
15
16use crate::wgsl_forge::ForgeError;
17
18/// The N-body step kernel source (embedded from the canonical `.wgsl`).
19pub const KIN_STEP_WGSL: &str = include_str!("../../shaders/kinematics.wgsl");
20/// Entry-point name of [`KIN_STEP_WGSL`].
21pub const KIN_STEP_ENTRY: &str = "nbody_step";
22/// Scalars per particle in the flat state buffer.
23pub const KIN_STRIDE: usize = 8;
24
25/// Exact CPU oracle for one N-body step. Reads `state_in`, returns the new state (same
26/// length), mirroring the WGSL scalar-for-scalar: forces accumulated in increasing-`j`
27/// order (skipping self), `1/(r²+soft)^{3/2}` via `r2 * sqrt(r2)`, then `v` then `x`
28/// using the new `v`. Particles with `mass == 0` take a zero inverse mass.
29pub fn nbody_step_cpu(state_in: &[f32], dt: f32, soft: f32, coupling: f32) -> Vec<f32> {
30    let count = state_in.len() / KIN_STRIDE;
31    let mut out = vec![0.0f32; state_in.len()];
32    for i in 0..count {
33        let bi = i * KIN_STRIDE;
34        let pix = state_in[bi];
35        let piy = state_in[bi + 1];
36        let piz = state_in[bi + 2];
37        let qi = state_in[bi + 7];
38
39        let mut fx = 0.0f32;
40        let mut fy = 0.0f32;
41        let mut fz = 0.0f32;
42        for j in 0..count {
43            if j == i {
44                continue;
45            }
46            let bj = j * KIN_STRIDE;
47            let rx = pix - state_in[bj];
48            let ry = piy - state_in[bj + 1];
49            let rz = piz - state_in[bj + 2];
50            let r2 = rx * rx + ry * ry + rz * rz + soft;
51            let inv = coupling * qi * state_in[bj + 7] / (r2 * r2.sqrt());
52            fx += rx * inv;
53            fy += ry * inv;
54            fz += rz * inv;
55        }
56
57        let mass = state_in[bi + 6];
58        let inv_m = if mass != 0.0 { 1.0 / mass } else { 0.0 };
59        let vx = state_in[bi + 3] + fx * inv_m * dt;
60        let vy = state_in[bi + 4] + fy * inv_m * dt;
61        let vz = state_in[bi + 5] + fz * inv_m * dt;
62
63        out[bi] = pix + vx * dt;
64        out[bi + 1] = piy + vy * dt;
65        out[bi + 2] = piz + vz * dt;
66        out[bi + 3] = vx;
67        out[bi + 4] = vy;
68        out[bi + 5] = vz;
69        out[bi + 6] = mass;
70        out[bi + 7] = qi;
71    }
72    out
73}
74
75/// Run one N-body step on the GPU and read back the new flat state. Builds a transient
76/// wgpu context, uploads `state_in` (binding 0, read), a zeroed output (binding 1,
77/// read_write) and `params = [dt, soft, coupling]` (binding 2, read), dispatches one
78/// invocation per particle, and reads the output back. Returns the same length as
79/// `state_in`.
80pub fn nbody_step_gpu(
81    state_in: &[f32],
82    dt: f32,
83    soft: f32,
84    coupling: f32,
85) -> Result<Vec<f32>, ForgeError> {
86    use crate::wgsl_forge::execute::{
87        BindingUsage, QualiaCompute, WgpuComputeContext, WgpuPipeline,
88    };
89    use crate::wgsl_forge::Schedule;
90
91    if state_in.is_empty() || state_in.len() % KIN_STRIDE != 0 {
92        return Err(ForgeError::GpuValidation(format!(
93            "nbody_step_gpu: state length {} is not a non-zero multiple of {KIN_STRIDE}",
94            state_in.len()
95        )));
96    }
97    let count = state_in.len() / KIN_STRIDE;
98    let capacity = (state_in.len() * 8).max(4 << 20);
99    let mut ctx = WgpuComputeContext::new(capacity)?;
100
101    let view_in = ctx.allocate_and_write(
102        bytemuck::cast_slice(state_in),
103        0,
104        0,
105        BindingUsage::StorageRead,
106    )?;
107    let zeros = vec![0.0f32; state_in.len()];
108    let view_out = ctx.allocate_and_write(
109        bytemuck::cast_slice(&zeros),
110        1,
111        0,
112        BindingUsage::StorageReadWrite,
113    )?;
114    let params = [dt, soft, coupling];
115    let view_params = ctx.allocate_and_write(
116        bytemuck::cast_slice(&params),
117        2,
118        0,
119        BindingUsage::StorageRead,
120    )?;
121
122    let buffers = vec![view_in, view_out, view_params];
123    let pipeline = WgpuPipeline::compile(&ctx, KIN_STEP_WGSL, KIN_STEP_ENTRY)?;
124    let schedule = Schedule {
125        workgroup_size: 64,
126        ..Default::default()
127    };
128    pipeline.dispatch(&buffers, &schedule, count)?;
129    let mut out = ctx.read_buffer_f32(&view_out)?;
130    out.truncate(state_in.len());
131    Ok(out)
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::wgsl_forge::validate::validate_wgsl;
138
139    /// The N-body kernel must naga-validate and expose the `nbody_step` entry point.
140    #[test]
141    fn kinematics_wgsl_validates() {
142        let report = validate_wgsl(KIN_STEP_WGSL).expect("kinematics WGSL must naga-validate");
143        assert!(
144            report.entry_points.iter().any(|e| e == KIN_STEP_ENTRY),
145            "validated module must expose {KIN_STEP_ENTRY}; got {:?}",
146            report.entry_points
147        );
148    }
149
150    /// Two equal positive charges on the x-axis must repel: with `coupling = +1`, the
151    /// left particle ends up moving in `−x` and the right in `+x`, so they separate.
152    #[test]
153    fn nbody_oracle_like_charges_repel() {
154        // p0 at x=0, p1 at x=1; both q=+1, m=1, at rest.
155        let state = vec![
156            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, // particle 0
157            1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, // particle 1
158        ];
159        let out = nbody_step_cpu(&state, 0.1, 1e-4, 1.0);
160        // Particle 0 (left) pushed −x → new x < 0; particle 1 (right) pushed +x → x > 1.
161        assert!(out[0] < 0.0, "left particle should move −x, got {}", out[0]);
162        assert!(
163            out[8] > 1.0,
164            "right particle should move +x, got {}",
165            out[8]
166        );
167        // x-velocities are equal and opposite (symmetric two-body).
168        assert!((out[3] + out[11]).abs() < 1e-5, "momentum not conserved");
169        // Mass/charge slots carried through unchanged.
170        assert_eq!(out[6], 1.0);
171        assert_eq!(out[7], 1.0);
172    }
173
174    /// Opposite charges attract: `q0=+1, q1=−1, coupling=+1` → they move together.
175    #[test]
176    fn nbody_oracle_opposite_charges_attract() {
177        let state = vec![
178            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, //
179            1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, -1.0, //
180        ];
181        let out = nbody_step_cpu(&state, 0.1, 1e-4, 1.0);
182        assert!(out[0] > 0.0, "left particle should move +x, got {}", out[0]);
183        assert!(
184            out[8] < 1.0,
185            "right particle should move −x, got {}",
186            out[8]
187        );
188    }
189
190    /// GPU certify: the kernel on a real adapter must match the CPU oracle within f32
191    /// tolerance over a deterministic multi-particle scene. Run by the orchestrator.
192    #[test]
193    #[serial_test::serial(gpu)]
194    fn nbody_gpu_matches_oracle() {
195        if !crate::wgsl_forge::test_gpu_available() {
196            return;
197        }
198        let count = 128usize;
199        let mut state = Vec::with_capacity(count * KIN_STRIDE);
200        for i in 0..count {
201            let f = i as f32;
202            state.extend_from_slice(&[
203                (f * 0.21) - 12.0,
204                (f * 0.13) - 8.0,
205                (f * 0.07) - 4.0,
206                0.0,
207                0.0,
208                0.0,
209                1.0 + (f % 3.0),
210                if i % 2 == 0 { 1.0 } else { -1.0 },
211            ]);
212        }
213        let (dt, soft, coupling) = (0.005f32, 1e-2, 1.0);
214        let expected = nbody_step_cpu(&state, dt, soft, coupling);
215        let gpu = nbody_step_gpu(&state, dt, soft, coupling).expect("nbody_step_gpu");
216        assert_eq!(gpu.len(), expected.len());
217        for (g, e) in gpu.iter().zip(expected.iter()) {
218            let tol = 1e-3 * e.abs().max(1.0);
219            assert!((g - e).abs() <= tol, "GPU/CPU mismatch: {g} vs {e}");
220        }
221    }
222}
qualia_core_db/wgsl_forge/physics/kinematics.rs

qualia_core_db/wgsl_forge/physics/
kinematics.rs
