Add GPU (wgpu) Direct Fourier transform and trajectory timing

Cargo.toml

@@ -26,6 +26,11 @@ serde = { version = "1.0", features = ["derive"] }
 yaml_serde = "0.10"
 voronota-ltr = { git = "https://github.com/mlund/voronota-ltr" }
 
+# GPU (wgpu) Direct Fourier transform — see `gpu` feature below.
+wgpu = { version = "29", optional = true }
+bytemuck = { version = "1", features = ["derive"], optional = true }
+pollster = { version = "0.4", optional = true }
+
 [dev-dependencies]
 approx = "0.5"
 criterion = "0.5"
@@ -34,6 +39,8 @@ tempfile = "3"
 [features]
 default = ["cli"]
 cli = ["dep:clap", "dep:pretty_env_logger", "dep:indicatif"]
+# GPU-accelerated Direct Fourier transform via wgpu/WGSL.
+gpu = ["dep:wgpu", "dep:bytemuck", "dep:pollster"]
 
 [[bin]]
 name = "pripps"

README.md

@@ -318,6 +318,9 @@ Requires a [Rust toolchain](https://www.rust-lang.org/tools/install):
 
 ```sh
 cargo install --git https://github.com/mlund/pripps.git
+
+# …or with GPU acceleration for the `direct` scheme (see below)
+cargo install --git https://github.com/mlund/pripps.git --features gpu
 ```
 
 ### Usage
@@ -369,6 +372,14 @@ With `--fit`, output is `q,i_exp,i_fit,sigma` at experimental
 q points. The experimental file is auto-detected as CSV or
 whitespace-separated (3 columns: q, I, σ).
 
+### GPU acceleration (`gpu` feature)
+
+Building with `--features gpu` runs the **`direct`** scheme on the GPU,
+which speeds up trajectory averaging substantially (measured ~15× over
+the CPU on an Apple M4). It applies to the vacuum transform; runs with a
+solvent model, or without a usable GPU, fall back to the CPU
+automatically. Pass `--cpu` to opt out.
+
 ### Environment variables
 
 | Variable | Description |

src/cli.rs

@@ -53,6 +53,11 @@ pub enum Commands {
         /// (the per-frame box from the XTC is used instead).
         #[clap(short = 'b', long = "box")]
         side_length: Option<f64>,
+        /// Force the CPU transform. When built with the `gpu` feature the
+        /// vacuum transform runs on the GPU by default; this opts back out.
+        /// (An active solvent model always uses the CPU.)
+        #[clap(long)]
+        cpu: bool,
         /// Solvent correction parameters (excluded volume and hydration)
         #[clap(flatten)]
         solvent: SolventArgs,
@@ -286,7 +291,9 @@ pub fn do_main() -> Result<()> {
                 let out = model.intensity(result.volume_scale, result.contrast_density);
                 crate::fit::write_fit_csv(&args.output, &exp, &out, result.scale)?;
             }
-            IntensityScheme::Direct { pmax } => {
+            // Fitting always uses the cached CPU DirectModel; the GPU path is
+            // for trajectory averaging, not the (c1, c2) optimisation loop.
+            IntensityScheme::Direct { pmax } | IntensityScheme::DirectGpu { pmax } => {
                 let box_sides = box_override.ok_or_else(|| {
                     crate::Error::usage("direct fitting requires box side lengths")
                 })?;
@@ -363,9 +370,16 @@ fn command_to_scheme(
         Commands::Direct {
             pmax,
             side_length,
+            cpu,
             solvent,
         } => Ok((
-            IntensityScheme::Direct { pmax },
+            // GPU is the default route when compiled in; `--cpu` forces the
+            // CPU path, and a build without the `gpu` feature is always CPU.
+            if cfg!(feature = "gpu") && !cpu {
+                IntensityScheme::DirectGpu { pmax }
+            } else {
+                IntensityScheme::Direct { pmax }
+            },
             side_length.map(|s| [s; 3]),
             solvent.to_config()?,
         )),

src/explicit.rs

@@ -280,7 +280,7 @@ fn eval_atom_form_factors(
 }
 
 /// Merge duplicate q-values by averaging, for 2-column data.
-fn average_duplicates(data: Vec<(f64, f64)>) -> Vec<(f64, f64)> {
+pub(crate) fn average_duplicates(data: Vec<(f64, f64)>) -> Vec<(f64, f64)> {
     let round = |x: f64| x.mul(1e6).round() as i64;
     let average = |pair: &[(f64, f64)]| {
         let mean_y = pair.iter().map(|(_, y)| y).sum::<f64>() / pair.len() as f64;