From 2f4cf4417535937f1486530e08e5631c6753e958 Mon Sep 17 00:00:00 2001 From: mlund Date: Thu, 18 Jun 2026 13:54:03 +0200 Subject: [PATCH 1/7] Add GPU (wgpu) Direct Fourier transform and trajectory timing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `gpu` cargo feature providing a wgpu/WGSL implementation of the Direct Fourier transform (vacuum path) for fast trajectory averaging. The kernel evaluates F(q) = Σ f_i·exp(i q·r_i), |F(q)|² in f32 with one workgroup per q-vector and a shared-memory reduction over atoms; static q-vectors and atom ids stay device-resident across frames. When built with the feature, `direct` runs on the GPU by default and `--cpu` forces the CPU path; it falls back to the CPU on an active solvent model or when no adapter is available. Degenerate sizes and device-side failures return errors rather than panicking. Measured ~15x over the rayon CPU path on an Apple M4 (88,400 atoms, 1300 q-vectors), matching the CPU f64 reference to <5e-6 relative. Also add offset-scan and per-frame timing logs to trajectory averaging. --- Cargo.toml | 7 + src/cli.rs | 18 +- src/explicit.rs | 2 +- src/gpu.rs | 546 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 49 +++++ src/trajectory.rs | 17 ++ 6 files changed, 636 insertions(+), 3 deletions(-) create mode 100644 src/gpu.rs diff --git a/Cargo.toml b/Cargo.toml index 9175bac..0aa1d8c 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/cli.rs b/src/cli.rs index 2438139..67bab6b 100644 --- a/src/cli.rs +++ b/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, + /// 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()?, )), diff --git a/src/explicit.rs b/src/explicit.rs index 051b494..4203508 100644 --- a/src/explicit.rs +++ b/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::() / pair.len() as f64; diff --git a/src/gpu.rs b/src/gpu.rs new file mode 100644 index 0000000..94cc60c --- /dev/null +++ b/src/gpu.rs @@ -0,0 +1,546 @@ +// Copyright 2025 Mikael Lund +// +// Licensed under the Apache license, version 2.0 (the "license"); +// you may not use this file except in compliance with the license. +// You may obtain a copy of the license at +// +// http://www.apache.org/licenses/license-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the license is distributed on an "as is" basis, +// without warranties or conditions of any kind, either express or implied. +// See the license for the specific language governing permissions and +// limitations under the license. + +//! GPU (wgpu/WGSL) implementation of the [`Direct`](crate::IntensityScheme::Direct) +//! Fourier transform, for the vacuum path only. +//! +//! The kernel evaluates, per reciprocal-lattice q-vector, +//! +//! ```text +//! F(q) = Σ_i f_i · exp(i q·r_i), I(q) = |F(q)|² +//! ``` +//! +//! exactly as the CPU [`DirectTransform`](crate::explicit) does, but on the +//! GPU in `f32`. The structure-factor sum has no catastrophic cancellation, so +//! `f32` keeps the relative error well below experimental SAXS noise (see +//! `WGPU-REPORT.md`); a golden-reference test pins it against the CPU path. +//! +//! The win is **trajectory averaging**: the q-vectors and atom ids stay +//! resident on the device across frames, and only the atom positions (and the +//! per-frame box / form-factor table) are re-uploaded each frame. + +use std::sync::Mutex; + +use wgpu::util::DeviceExt; + +use crate::explicit::{average_duplicates, directions}; +use crate::formfactor::FormFactor; +use crate::{Error, FormFactorMap, Intensity, IntensityCalculator, Result, Structure, Vector3}; + +/// One workgroup reduces over all atoms for a single q-vector. The thread +/// count per workgroup is fixed at 256 in the WGSL `@workgroup_size` and the +/// matching shared-memory array sizes and reduction stride below. +const SHADER: &str = r#" +struct Params { + inv_box: vec4, // 1/box_sides; .xyz used + n_atoms: u32, + n_q: u32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var params: Params; +@group(0) @binding(1) var directions: array>; // n_q (.xyz used) +@group(0) @binding(2) var positions: array>; // n_atoms (.xyz used) +@group(0) @binding(3) var atom_ids: array; // n_atoms +@group(0) @binding(4) var ff_table: array; // n_kinds * n_q (kind-major) +@group(0) @binding(5) var out_iq: array; // n_q + +var sh_re: array; +var sh_im: array; + +@compute @workgroup_size(256) +fn structure_factor( + @builtin(workgroup_id) wg: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let qi = wg.x; + let q = directions[qi].xyz * params.inv_box.xyz; + + // Strided walk over atoms; each thread accumulates its own partial sum. + var re = 0.0; + var im = 0.0; + var a = lid.x; + loop { + if (a >= params.n_atoms) { break; } + let r = positions[a].xyz; + let f = ff_table[atom_ids[a] * params.n_q + qi]; + let qr = dot(q, r); + re = re + f * cos(qr); + im = im + f * sin(qr); + a = a + 256u; + } + sh_re[lid.x] = re; + sh_im[lid.x] = im; + workgroupBarrier(); + + // Shared-memory tree reduction. + var stride = 128u; + loop { + if (stride == 0u) { break; } + if (lid.x < stride) { + sh_re[lid.x] = sh_re[lid.x] + sh_re[lid.x + stride]; + sh_im[lid.x] = sh_im[lid.x] + sh_im[lid.x + stride]; + } + workgroupBarrier(); + stride = stride / 2u; + } + + if (lid.x == 0u) { + let rr = sh_re[0]; + let ii = sh_im[0]; + out_iq[qi] = rr * rr + ii * ii; + } +} +"#; + +/// Per-frame uniform block. `#[repr(C)]` with explicit padding so the layout +/// matches the WGSL `Params` struct (vec4 + two u32 + two pad u32 = 32 bytes, +/// the size rounding required by the vec4's 16-byte alignment). +#[repr(C)] +#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] +struct Params { + inv_box: [f32; 4], + n_atoms: u32, + n_q: u32, + _pad: [u32; 2], +} + +/// GPU device context plus the static, structure-independent resources +/// (q-directions, pipeline). Built once in [`GpuDirectTransform::new`]. +struct GpuContext { + device: wgpu::Device, + queue: wgpu::Queue, + pipeline: wgpu::ComputePipeline, + layout: wgpu::BindGroupLayout, + /// q-directions as `vec4` (xyz used); divided by the box in-shader. + dir_buf: wgpu::Buffer, + n_q: u32, +} + +/// Structure-sized resources, (re)built when the atom count first becomes +/// known or changes. Positions, the form-factor table, and the box are +/// rewritten every frame; ids stay fixed across a trajectory. +struct FrameState { + n_atoms: u32, + pos_buf: wgpu::Buffer, + ff_buf: wgpu::Buffer, + params_buf: wgpu::Buffer, + out_buf: wgpu::Buffer, + read_buf: wgpu::Buffer, + bind_group: wgpu::BindGroup, +} + +/// GPU Direct Fourier transform (vacuum path). Implements +/// [`IntensityCalculator`] and slots in alongside the CPU +/// [`DirectTransform`](crate::explicit::DirectTransform). +pub(crate) struct GpuDirectTransform { + formfactors: FormFactorMap, + /// q-directions in host (f64) form, for per-frame `|q|` and box division. + directions: Vec, + ctx: GpuContext, + /// Lazily built on the first frame; the trait is `&self`, so interior + /// mutability is required. The trajectory loop is serial, so contention + /// is nil. + frame: Mutex>, +} + +impl GpuDirectTransform { + /// Initialise wgpu and compile the kernel. Returns `Err` if no suitable + /// GPU adapter is available, so the caller can fall back to the CPU path. + pub(crate) fn new(formfactors: FormFactorMap, pmax: usize) -> Result { + let directions = directions(pmax); + let n_q = directions.len() as u32; + + let instance = wgpu::Instance::default(); + let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::HighPerformance, + ..Default::default() + })) + .map_err(|e| Error::internal(format!("no GPU adapter available: {e}")))?; + + log::info!("GPU adapter: {}", adapter.get_info().name); + + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("pripps-direct"), + ..Default::default() + })) + .map_err(|e| Error::internal(format!("failed to create GPU device: {e}")))?; + + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("direct-structure-factor"), + source: wgpu::ShaderSource::Wgsl(SHADER.into()), + }); + + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("direct-bind-layout"), + entries: &bind_layout_entries(), + }); + + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("direct-pipeline-layout"), + bind_group_layouts: &[Some(&layout)], + immediate_size: 0, + }); + + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("direct-pipeline"), + layout: Some(&pipeline_layout), + module: &module, + entry_point: Some("structure_factor"), + compilation_options: Default::default(), + cache: None, + }); + + // Static q-direction buffer (xyz packed into vec4 for alignment). + let dir_data: Vec<[f32; 4]> = directions + .iter() + .map(|d| [d.x as f32, d.y as f32, d.z as f32, 0.0]) + .collect(); + let dir_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("directions"), + contents: bytemuck::cast_slice(&dir_data), + usage: wgpu::BufferUsages::STORAGE, + }); + + Ok(Self { + formfactors, + directions, + ctx: GpuContext { + device, + queue, + pipeline, + layout, + dir_buf, + n_q, + }, + frame: Mutex::new(None), + }) + } +} + +impl IntensityCalculator for GpuDirectTransform { + fn calc_intensities(&self, structure: &Structure) -> Result { + let box_sides = structure.box_sides.ok_or_else(|| { + Error::usage("direct Fourier transform requires box side lengths on the structure") + })?; + + // |q| for every direction; a pure function of the (per-frame) box. + let qmags: Vec = self + .directions + .iter() + .map(|d| d.component_div(&box_sides).norm()) + .collect(); + + // Degenerate sizes would create zero-length GPU buffers (a hard wgpu + // validation error), so handle them on the host exactly as the CPU + // path does: no q-vectors → empty curve; no atoms → all-zero I(q). + if qmags.is_empty() { + return Ok(Intensity::default()); + } + if structure.pos.is_empty() { + let table = qmags.into_iter().map(|q| (q, 0.0)).collect(); + let (q, total): (Vec, Vec) = average_duplicates(table).into_iter().unzip(); + return Ok(Intensity { + q, + total, + ..Default::default() + }); + } + + let n_atoms = structure.pos.len() as u32; + let n_kinds = self.formfactors.len() as u32; + let ctx = &self.ctx; + + // (Re)build the structure-sized buffers if this is the first frame or + // the atom count changed (it doesn't within a trajectory). + let mut guard = self.frame.lock().map_err(|_| { + Error::internal("GPU frame state mutex was poisoned by an earlier panic") + })?; + if guard.as_ref().map(|f| f.n_atoms) != Some(n_atoms) { + *guard = Some(ctx.build_frame_state(structure, n_kinds)); + } + let frame = guard + .as_ref() + .ok_or_else(|| Error::internal("GPU frame state missing after initialization"))?; + + // Per-frame host work: form-factor table, packed positions, box. + let ff_table = self.form_factor_table(&qmags)?; + + let positions: Vec<[f32; 4]> = structure + .pos + .iter() + .map(|p| [p.x as f32, p.y as f32, p.z as f32, 0.0]) + .collect(); + + let params = Params { + inv_box: [ + (1.0 / box_sides.x) as f32, + (1.0 / box_sides.y) as f32, + (1.0 / box_sides.z) as f32, + 0.0, + ], + n_atoms, + n_q: ctx.n_q, + _pad: [0, 0], + }; + + ctx.queue + .write_buffer(&frame.pos_buf, 0, bytemuck::cast_slice(&positions)); + ctx.queue + .write_buffer(&frame.ff_buf, 0, bytemuck::cast_slice(&ff_table)); + ctx.queue + .write_buffer(&frame.params_buf, 0, bytemuck::bytes_of(¶ms)); + + let intensities = ctx.dispatch(frame)?; + + // Pair each |q| with its I(q) and average duplicate magnitudes, exactly + // as the CPU vacuum path does. + let table: Vec<(f64, f64)> = qmags + .into_iter() + .zip(intensities.into_iter().map(f64::from)) + .collect(); + let (q, total): (Vec, Vec) = average_duplicates(table).into_iter().unzip(); + + Ok(Intensity { + q, + total, + ..Default::default() + }) + } +} + +impl GpuDirectTransform { + /// Build the `[n_kinds × n_q]` (kind-major) form-factor table in `f32`, + /// evaluating each kind's form factor at every q-magnitude. + fn form_factor_table(&self, qmags: &[f64]) -> Result> { + let ffs: Vec<&FormFactor> = self.formfactors.form_factors().collect(); + let mut table = Vec::with_capacity(ffs.len() * qmags.len()); + for ff in &ffs { + for &q in qmags { + table.push(ff.eval(q)? as f32); + } + } + Ok(table) + } +} + +impl GpuContext { + /// Allocate the structure-sized buffers and the bind group. `atom_ids` is + /// uploaded once here (constant across a trajectory); positions, the + /// form-factor table, and the box are written per frame. + fn build_frame_state(&self, structure: &Structure, n_kinds: u32) -> FrameState { + let n_atoms = structure.pos.len() as u32; + let device = &self.device; + + let ids: Vec = structure.ids.iter().map(|&id| id as u32).collect(); + let ids_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("atom-ids"), + contents: bytemuck::cast_slice(&ids), + usage: wgpu::BufferUsages::STORAGE, + }); + + let pos_buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("positions"), + size: u64::from(n_atoms) * 16, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let ff_buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("ff-table"), + size: u64::from(n_kinds) * u64::from(self.n_q) * 4, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let params_buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("params"), + size: std::mem::size_of::() as u64, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let out_size = u64::from(self.n_q) * 4; + let out_buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("out-iq"), + size: out_size, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let read_buf = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("readback"), + size: out_size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + + let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("direct-bind-group"), + layout: &self.layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: params_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: self.dir_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: pos_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: ids_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 4, + resource: ff_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 5, + resource: out_buf.as_entire_binding(), + }, + ], + }); + + FrameState { + n_atoms, + pos_buf, + ff_buf, + params_buf, + out_buf, + read_buf, + bind_group, + } + } + + /// Dispatch one workgroup per q-vector, copy the result to the readback + /// buffer, and block until it is host-visible. Returns I(q) per q-vector. + /// + /// Any device-side failure (lost device, failed map) is returned as an + /// [`Error`] rather than panicking, so a transient GPU fault mid-trajectory + /// surfaces as a recoverable error instead of aborting the process. + fn dispatch(&self, frame: &FrameState) -> Result> { + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("direct-encoder"), + }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("direct-pass"), + timestamp_writes: None, + }); + pass.set_pipeline(&self.pipeline); + pass.set_bind_group(0, &frame.bind_group, &[]); + pass.dispatch_workgroups(self.n_q, 1, 1); + } + let out_size = u64::from(self.n_q) * 4; + encoder.copy_buffer_to_buffer(&frame.out_buf, 0, &frame.read_buf, 0, out_size); + self.queue.submit([encoder.finish()]); + + let slice = frame.read_buf.slice(..); + let (tx, rx) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |res| { + let _ = tx.send(res); + }); + self.device + .poll(wgpu::PollType::wait_indefinitely()) + .map_err(|e| Error::internal(format!("GPU poll failed: {e}")))?; + rx.recv() + .map_err(|_| Error::internal("GPU readback callback was dropped"))? + .map_err(|e| Error::internal(format!("GPU buffer map failed: {e}")))?; + + let data = slice.get_mapped_range(); + let out = bytemuck::cast_slice::(&data).to_vec(); + drop(data); + frame.read_buf.unmap(); + Ok(out) + } +} + +/// Bind-group layout: one uniform (params), four read-only storage inputs, +/// one read-write storage output. +fn bind_layout_entries() -> [wgpu::BindGroupLayoutEntry; 6] { + let storage = |binding: u32, read_only: bool| wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }; + [ + wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }, + storage(1, true), + storage(2, true), + storage(3, true), + storage(4, true), + storage(5, false), + ] +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Pripps; + use crate::explicit::DirectTransform; + + /// The GPU (f32) vacuum transform must agree with the CPU (f64) + /// `DirectTransform` to within SAXS-grade tolerance. Skips gracefully when + /// no GPU adapter is available (e.g. headless CI), so the suite still passes. + #[test] + fn gpu_matches_cpu_direct_vacuum() { + let yaml = std::path::Path::new("assets/poly_amino_acid.yaml"); + let xyz = std::path::Path::new("tests/lysozyme/4lzt.xyz"); + let pripps = Pripps::from_yaml(yaml).unwrap(); + + let gpu = match GpuDirectTransform::new(pripps.formfactors.clone(), 4) { + Ok(g) => g, + Err(e) => { + eprintln!("skipping GPU golden-reference test: {e}"); + return; + } + }; + let cpu = DirectTransform::new(pripps.formfactors.clone(), 4, None); + + let mut structure = Structure::from_file(xyz, &pripps.atomkinds()).unwrap(); + structure.box_sides = Some(Vector3::new(200.0, 200.0, 200.0)); + + let cpu_i = cpu.calc_intensities(&structure).unwrap(); + let gpu_i = gpu.calc_intensities(&structure).unwrap(); + + assert_eq!(cpu_i.q, gpu_i.q, "q grids must match exactly"); + for i in 0..cpu_i.total.len() { + let (a, b) = (cpu_i.total[i], gpu_i.total[i]); + let rel = (a - b).abs() / a.abs().max(1e-30); + assert!( + rel < 1e-3, + "q={}: cpu={a} gpu={b} rel_err={rel}", + cpu_i.q[i] + ); + } + } +} diff --git a/src/lib.rs b/src/lib.rs index baac710..0d7783b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,6 +25,8 @@ mod error; mod explicit; mod fit; mod formfactor; +#[cfg(feature = "gpu")] +mod gpu; mod input; mod multipole; mod q_sampler; @@ -114,6 +116,15 @@ pub enum IntensityScheme { /// Maximum Miller-index multiple used to build q-vectors. pmax: usize, }, + /// GPU (wgpu) Direct Fourier transform, vacuum path only. Identical + /// physics to [`Direct`](Self::Direct) but evaluated on the GPU in `f32`, + /// for fast trajectory averaging. Requires the `gpu` build feature; with + /// an active solvent or no available adapter it falls back to the CPU + /// [`Direct`](Self::Direct) path. + DirectGpu { + /// Maximum Miller-index multiple used to build q-vectors. + pmax: usize, + }, /// Spherical-harmonic multipole expansion /// (*Acta Cryst. A51*, 1995). O(N·Q·L²), linear in atom count — /// preferred for large structures and iterative fitting. `l_max` @@ -332,6 +343,44 @@ impl Pripps { active_solvent.cloned(), )) } + IntensityScheme::DirectGpu { pmax } => { + #[cfg(feature = "gpu")] + { + // The GPU path covers the vacuum case only; with an active + // solvent or no usable adapter, fall back to the CPU Direct + // transform so results stay correct. + if active_solvent.is_some() { + log::warn!( + "solvent corrections are not supported on the GPU; using CPU direct transform" + ); + Box::new(DirectTransform::new( + self.formfactors.clone(), + pmax, + active_solvent.cloned(), + )) + } else { + match gpu::GpuDirectTransform::new(self.formfactors.clone(), pmax) { + Ok(calc) => { + log::info!( + "Using GPU direct Fourier transform with p_max = {pmax}" + ); + Box::new(calc) as Box + } + Err(e) => { + log::warn!("GPU unavailable ({e}); using CPU direct transform"); + Box::new(DirectTransform::new(self.formfactors.clone(), pmax, None)) + } + } + } + } + #[cfg(not(feature = "gpu"))] + { + let _ = pmax; + return Err(Error::usage( + "GPU direct transform requested but pripps was built without the `gpu` feature", + )); + } + } IntensityScheme::Multipole { qmin, qmax, diff --git a/src/trajectory.rs b/src/trajectory.rs index f72863f..17a1136 100644 --- a/src/trajectory.rs +++ b/src/trajectory.rs @@ -16,6 +16,7 @@ use itertools::Itertools; use molly::{Frame, XTCReader, selection::AtomSelection}; use std::iter::zip; use std::num::NonZeroUsize; +use std::time::Instant; use std::{f64, path}; #[cfg(feature = "cli")] @@ -164,8 +165,17 @@ pub fn trajectory_average( structure: &Structure, xtc_opts: crate::TrajectoryOptions, ) -> Result { + // The offset scan walks every frame header in the file (see + // `XtcIterator::open`); for multi-gigabyte trajectories this is a + // noticeable up-front cost, so time it separately from the averaging. + let scan_start = Instant::now(); let reader = XtcIterator::open(xtc_opts.xtcfile, xtc_opts.weights, xtc_opts.stride)?; let num_frames = reader.len(); + log::info!( + "Scanned frame offsets in {:.2} s", + scan_start.elapsed().as_secs_f64() + ); + #[cfg(feature = "cli")] let progress = ProgressBar::new(num_frames as u64); @@ -175,6 +185,7 @@ pub fn trajectory_average( let mut structure = structure.clone(); let mut weight_sum = 0.0; + let average_start = Instant::now(); for (frame, weight) in reader { weight_sum += weight; structure.set_positions(&frame)?; @@ -193,6 +204,12 @@ pub fn trajectory_average( progress.inc(1); } + let elapsed = average_start.elapsed().as_secs_f64(); + log::info!( + "Averaged {num_frames} frames in {elapsed:.2} s ({:.1} ms/frame)", + 1e3 * elapsed / num_frames.max(1) as f64 + ); + // Normalize the accumulated sum by the total weight to get the weighted average for acc in &mut mean.total { *acc /= weight_sum; From 5cb9b45d1dd36f1bb79e6c99e4922f0c1c033c4e Mon Sep 17 00:00:00 2001 From: mlund Date: Thu, 18 Jun 2026 14:36:52 +0200 Subject: [PATCH 2/7] Address PR #13 review: count processed frames; note GPU precision split Time trajectory averaging against the number of frames actually consumed rather than the offset-table length, which overstates frames (and skews ms/frame) if XtcIterator stops early on a frame read error. Document why the GPU q-magnitudes stay f64 (form factors and q-grid match the CPU path; only the in-shader phase is f32). --- src/gpu.rs | 2 ++ src/trajectory.rs | 10 ++++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/gpu.rs b/src/gpu.rs index 94cc60c..c73aa4d 100644 --- a/src/gpu.rs +++ b/src/gpu.rs @@ -237,6 +237,8 @@ impl IntensityCalculator for GpuDirectTransform { })?; // |q| for every direction; a pure function of the (per-frame) box. + // Kept f64 to match the CPU form factors and q-grid; only the in-shader + // phase `q·r` is f32. let qmags: Vec = self .directions .iter() diff --git a/src/trajectory.rs b/src/trajectory.rs index 17a1136..c4d10dc 100644 --- a/src/trajectory.rs +++ b/src/trajectory.rs @@ -185,6 +185,11 @@ pub fn trajectory_average( let mut structure = structure.clone(); let mut weight_sum = 0.0; + // Count frames we actually consume: `XtcIterator` stops early on a frame + // read error (its `next` swallows the error), so `num_frames` (the offset + // table length) can overstate what was processed. Time against the real + // count so ms/frame stays honest. + let mut processed = 0usize; let average_start = Instant::now(); for (frame, weight) in reader { weight_sum += weight; @@ -200,14 +205,15 @@ pub fn trajectory_average( *acc += weight * v; } + processed += 1; #[cfg(feature = "cli")] progress.inc(1); } let elapsed = average_start.elapsed().as_secs_f64(); log::info!( - "Averaged {num_frames} frames in {elapsed:.2} s ({:.1} ms/frame)", - 1e3 * elapsed / num_frames.max(1) as f64 + "Averaged {processed} frames in {elapsed:.2} s ({:.1} ms/frame)", + 1e3 * elapsed / processed.max(1) as f64 ); // Normalize the accumulated sum by the total weight to get the weighted average From 9ca866bbcff13a3573cb7266a7d0e3fd7681e980 Mon Sep 17 00:00:00 2001 From: mlund Date: Thu, 18 Jun 2026 14:36:52 +0200 Subject: [PATCH 3/7] docs: document the gpu feature in the README Note the optional `gpu` build feature: GPU-accelerated `direct` scheme for trajectory averaging, automatic CPU fallback, and the `--cpu` opt-out. --- README.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/README.md b/README.md index 17e2594..05441db 100644 --- a/README.md +++ b/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 | From e5dcfc2fbe193f10ea4fe51d58ca7b97e9a91873 Mon Sep 17 00:00:00 2001 From: mlund Date: Thu, 18 Jun 2026 18:21:06 +0200 Subject: [PATCH 4/7] Merge multiple --atomfile files; split gaussian_atoms.yaml Form-factor definitions can now be composed from several YAML files. Repeat --atomfile (or pass several paths), or give the Python Pripps constructor a list. Files are merged in order; on a duplicate species the last file wins and a warning is logged. Merging happens before symbolic resolution, so an !Alias/!United/!Linear in one file may reference a species defined in another. - formfactor.rs: split read_formfactors into parse_specs + finalize_specs; add read_formfactors_merged. - lib.rs: add Pripps::from_yaml_files; from_yaml delegates via from_formfactors. - cli.rs: --atomfile becomes Vec (repeatable, order-preserving). - pripps-py: Pripps(...) accepts a str or list[str]. Split gaussian_atoms.yaml into an atomic/united-atom-only table: drop the 20 coarse-grained residue beads (byte-identical duplicates of poly_amino_acid.yaml) and 8 unused residue-level aliases, and rewrite the header. Residue-bead consumers load poly_amino_acid.yaml (or merge it with the atomic file). This resolves the mixed-representation file and supersedes the rename-only approach in #11. Closes #7 --- README.md | 26 ++- assets/gaussian_atoms.yaml | 376 +------------------------------------ pripps-py/src/lib.rs | 23 ++- src/cli.rs | 10 +- src/formfactor.rs | 111 ++++++++++- src/lib.rs | 23 ++- 6 files changed, 180 insertions(+), 389 deletions(-) diff --git a/README.md b/README.md index f89b30d..38fa356 100644 --- a/README.md +++ b/README.md @@ -234,12 +234,33 @@ scaled by `volume_scale` — instead of building a Gaussian dummy from `radius`/`displaced_volume`. Every species present in the structure must define it when `--excluded per-kind` is selected. +#### Merging multiple files + +Several form-factor files can be supplied and are merged in the order +given — repeat `--atomfile` (or pass several paths) on the CLI, or pass a +list to the Python constructor: + +```bash +pripps --atomfile assets/gaussian_atoms.yaml --atomfile my_overrides.yaml ... +``` + +```python +pr = pripps.Pripps(["assets/gaussian_atoms.yaml", "my_overrides.yaml"]) +``` + +This lets you compose, e.g., a coarse-grained residue file with an atomic +file. **Duplicate policy:** if the same species is defined in more than one +file, the **last file listed wins** and a warning is logged naming the +overridden species. Merging happens before symbolic resolution, so an +`!Alias` / `!United` / `!Linear` entry in one file may reference a species +defined in another. + ### Shipped form-factor files | File | Description | Best for | |------|-------------|----------| | `assets/foxs_formfactors.yaml` | Cromer-Mann (Int. Tables Vol. C) atomic form factors, FoXS-compatible volumes and SASA radii | Atomic PDBs with `--excluded foxs` | -| `assets/gaussian_atoms.yaml` | Mixed table: residue single-bead polynomials plus atomic/united-atom Gaussians (CH/CH₂/NH/…) | Atomic PDBs with PepsiSAXS-style solvent models | +| `assets/gaussian_atoms.yaml` | Atomic and united-atom Gaussians (CH/CH₂/NH/…), PepsiSAXS-style (K = 4π) | Atomic PDBs with PepsiSAXS-style solvent models | | `assets/stovgaard2010.yaml` | Per-residue form-factor centroids (one-body model, q ∈ [0, 0.75] Å⁻¹) from Stovgaard et al. 2010, | Residue / Cα single-bead models | | `assets/poly_amino_acid.yaml` | Amino-acid bead polynomials with a separate `excluded_solvent` polynomial per bead | Cα / single-bead with `--excluded per-kind` | | `assets/poly_amino_acid_combined.yaml` | Amino-acid bead polynomials with the excluded solvent already folded into the form factor | Cα / single-bead with `--excluded disabled` | @@ -335,6 +356,9 @@ pripps --help # Debye formula (exact, O(N²)) pripps -a atoms.yaml -i structure.pdb debye +# Merge several form-factor files (last listed wins on duplicate species) +pripps -a base.yaml -a overrides.yaml -i structure.pdb debye + # Multipole expansion (linear in atoms, truncation auto-selected) pripps -a atoms.yaml -i structure.pdb multipole pripps -a atoms.yaml -i structure.pdb multipole --lmax 68 # fixed order diff --git a/assets/gaussian_atoms.yaml b/assets/gaussian_atoms.yaml index 5f072aa..1da56a2 100644 --- a/assets/gaussian_atoms.yaml +++ b/assets/gaussian_atoms.yaml @@ -1,364 +1,8 @@ -# Amino acid single bead approximation for atomistic scattering -ALA: - { - radius: 3.1, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 37.98651145079954, - 0.0, - -17.88881261159881, - -3.1097731825829333, - 8.846459092736989, - -2.3572911188270753, - 0.05331786541354333, - ], - }, - } -ARG: - { - radius: 4.0, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 83.96908927242093, - 0.0, - -278.81910994314467, - 490.5059319443547, - -383.9755997417938, - 143.65451627028935, - -20.783481179736413, - ], - }, - } -ASN: - { - radius: 3.6, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 59.98205560472268, - 0.0, - -40.54214890095476, - -32.39953367060282, - 73.4159357342896, - -37.17064504604371, - 6.101333446616579, - ], - }, - } -ASP: - { - radius: 3.6, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 59.98252699985404, - 0.0, - -37.42179690845112, - -37.97031252961563, - 77.22160318216388, - -38.25436485287656, - 6.198643717045879, - ], - }, - } -CYS: - { - radius: 3.6, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 53.98471591625841, - 0.0, - -28.55463591537983, - -34.95759123047967, - 63.496177867671385, - -29.839780719449596, - 4.609281628955767, - ], - }, - } -GLN: - { - radius: 3.8, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 68.00029704612041, - 0.0, - -97.52724702286758, - 63.863879205326334, - 13.865971310987792, - -22.96589467520148, - 5.296934807386833, - ], - }, - } -GLU: - { - radius: 3.8, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 68.00262098772043, - 0.0, - -95.50455833177392, - 59.449715597229726, - 18.625430344038303, - -25.334425917186397, - 5.722041637792017, - ], - }, - } -GLY: - { - radius: 2.9, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 29.99006405529222, - 0.0, - -11.572242769967232, - -0.9021293993630419, - 3.790951509543525, - -0.7238332866526376, - -0.05549468364095844, - ], - }, - } -HIS: - { - radius: 3.9, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 71.99289906060545, - 0.0, - -93.58672742143521, - 41.460338892548805, - 39.724339660983915, - -34.41529003289915, - 7.057596750057414, - ], - }, - } -ILE: - { - radius: 3.6, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 61.97933054508859, - 0.0, - -45.55749645626756, - -26.94147636420316, - 68.41091746940783, - -34.307325839550984, - 5.508363044131372, - ], - }, - } -LEU: - { - radius: 3.6, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 61.9874112464829, - 0.0, - -54.07829774966109, - -19.672724727921118, - 75.75933357599484, - -42.87909415253206, - 7.580643479500708, - ], - }, - } -LYS: - { - radius: 3.7, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 69.99129116923348, - 0.0, - -177.94423757968704, - 272.0256856055207, - -188.2920973172914, - 63.303038257186586, - -8.32612812430321, - ], - }, - } -MET: - { - radius: 3.8, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 70.00691278772638, - 0.0, - -104.22345654704631, - 73.82601789465483, - 10.126507627913185, - -23.351946238697035, - 5.598503080281802, - ], - }, - } -PHE: - { - radius: 3.9, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 77.99863587497389, - 0.0, - -135.30506501839736, - 122.61728717372222, - -26.829761151876106, - -9.267313845762843, - 3.479368368324131, - ], - }, - } -PRO: - { - radius: 3.4, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 52.98008056903471, - 0.0, - -28.025579140605455, - -17.945632640349288, - 36.70484993598286, - -16.388585233440434, - 2.3759607129713807, - ], - }, - } -SER: - { - radius: 3.3, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 45.98433274188738, - 0.0, - -23.004358312306444, - -13.52528557404974, - 25.318714967455264, - -9.60754708227119, - 1.0838180112873488, - ], - }, - } -THR: - { - radius: 3.5, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 53.98079904257712, - 0.0, - -30.205614210054552, - -20.68857584136233, - 39.80321996161699, - -16.614363714862087, - 2.188488144363049, - ], - }, - } -TRP: - { - radius: 4.3, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 98.00026950327825, - 0.0, - -239.82185506986724, - 306.5877727946093, - -164.7189974469082, - 39.67855896479587, - -3.263270784166963, - ], - }, - } -TYR: - { - radius: 4.1, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 85.99025648157047, - 0.0, - -205.8271529828311, - 282.4197901545501, - -174.14171820412088, - 52.3398359916699, - -6.16839805508144, - ], - }, - } -VAL: - { - radius: 3.4, - formfactor: - !Polynomial { - qrange: [0.0, 0.75], - coeff: - [ - 53.97802539665846, - 0.0, - -30.730380715977926, - -23.1259991191962, - 43.982242401804015, - -18.76619704074442, - 2.5506887609142446, - ], - }, - } +# Atomic / united-atom form factors for PRIPPS (PepsiSAXS-style K = 4π). +# +# Atomic-resolution structures only. Residue single-bead form factors live +# in assets/poly_amino_acid.yaml — pass several --atomfile paths to merge +# this atomic table with a coarse-grained residue file when needed. # Atomic properties for PRIPPS # @@ -810,17 +454,7 @@ I: CTR: { radius: 2.0, formfactor: !Alias C } NTR: { radius: 2.0, formfactor: !Alias N } HNTR: { ff: !Alias N } -HLYS: { ff: !Alias LYS } -HARG: { ff: !Alias ARG } -HHIS: { ff: !Alias HIS } -HGLU: { ff: !Alias GLU } -HASP: { ff: !Alias ASP } HCTR: { ff: !Alias C } -HLYS: { ff: !Alias LYS } -HLYS: { ff: !Alias LYS } -HTYR: { ff: !Alias TYR } -HCYS: { ff: !Alias CYS } -CYX: { ff: !Alias CYS } # PDB atom name aliases — map multi-character PDB names to their # element form factors so atomic-resolution structures work with diff --git a/pripps-py/src/lib.rs b/pripps-py/src/lib.rs index 0b7a728..aecfda5 100644 --- a/pripps-py/src/lib.rs +++ b/pripps-py/src/lib.rs @@ -259,7 +259,8 @@ impl PyFitResult { /// Scattering calculator loaded from a form-factor YAML file. /// -/// Construct with `Pripps(atomfile)`, then either build a cached model +/// Construct with `Pripps(atomfile)` — a single path or a list of paths +/// merged in order (latest wins on duplicate species) — then either build a cached model /// (`multipole_model(...)` / `debye_model(...)` / `direct_model(...)`) /// for repeated/fitting use, or call `intensity(...)` for a one-shot /// calculation (the only path that averages over an XTC trajectory). @@ -268,10 +269,24 @@ struct PyPripps(Pripps); #[pymethods] impl PyPripps { - /// Load a pripps instance from a form-factor YAML file. + /// Load a pripps instance from one or more form-factor YAML files. + /// + /// `atomfile` is either a single path (`Pripps("a.yaml")`) or a list + /// of paths (`Pripps(["base.yaml", "override.yaml"])`). Multiple files + /// are merged in order; when a species appears in more than one file + /// the last one listed wins (a warning is logged for each override). #[new] - fn new(atomfile: &str) -> PyResult { - Pripps::from_yaml(atomfile.as_ref()) + fn new(atomfile: &Bound<'_, PyAny>) -> PyResult { + let paths: Vec = if let Ok(single) = atomfile.extract::() { + vec![single.into()] + } else if let Ok(many) = atomfile.extract::>() { + many.into_iter().map(PathBuf::from).collect() + } else { + return Err(PyTypeError::new_err( + "`atomfile` must be a str or a list of str", + )); + }; + Pripps::from_yaml_files(&paths) .map(Self) .map_err(into_py_err) } diff --git a/src/cli.rs b/src/cli.rs index 67bab6b..eb5391b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -216,9 +216,11 @@ pub struct Args { #[clap(subcommand)] pub command: Commands, - /// Form factor definitions for species (atoms, residues, etc.) - #[clap(long, short = 'a')] - pub atomfile: PathBuf, + /// Form factor definitions for species (atoms, residues, etc.). + /// Repeat the flag or pass several paths to merge multiple files; + /// on a duplicate species the last file listed wins. + #[clap(long = "atomfile", short = 'a', num_args = 1.., required = true)] + pub atomfiles: Vec, /// Input structure (xyz, pdb, gro etc.) #[clap(short = 'i', long)] @@ -266,7 +268,7 @@ pub fn do_main() -> Result<()> { } pretty_env_logger::init(); - let pripps = Pripps::from_yaml(&args.atomfile)?; + let pripps = Pripps::from_yaml_files(&args.atomfiles)?; let (scheme, box_override, solvent) = command_to_scheme(args.command)?; if let Some(fit_path) = &args.fit { diff --git a/src/formfactor.rs b/src/formfactor.rs index 2eaf897..4a85305 100644 --- a/src/formfactor.rs +++ b/src/formfactor.rs @@ -18,7 +18,13 @@ use crate::{AtomKind, Error, Result}; use itertools::Itertools; use nalgebra::{DMatrix, DVector}; use serde::{Deserialize, Serialize}; -use std::{cmp::Ordering, collections::BTreeMap, fs::File, ops::Add, path::Path}; +use std::{ + cmp::Ordering, + collections::BTreeMap, + fs::File, + ops::Add, + path::{Path, PathBuf}, +}; /// A flat table pairing each [`AtomKind`] with its runtime [`FormFactor`], /// stored in atom-kind id order. Lookups by id are `O(1)` vector indexes. @@ -401,11 +407,19 @@ fn interpolate(table: &[(f64, f64)], x: f64) -> f64 { y0 + (x - x0) * (y1 - y0) / (x1 - x0) } -/// Read YAML file with form factors for different atom types (atomic or coarse grained) -pub fn read_formfactors(path: &Path) -> Result { +/// Parse a YAML form-factor file into its raw, still-symbolic spec map +/// (keyed by species name). The symbolic variants (`!Alias`, `!United`, +/// `!Linear`) are left unresolved so that [`read_formfactors_merged`] can +/// combine maps from several files *before* resolution — letting an alias +/// in one file reference a definition supplied by another. +fn parse_specs(path: &Path) -> Result> { let file = File::open(path).with_path(path)?; - let mut specs: BTreeMap = yaml_serde::from_reader(file)?; + Ok(yaml_serde::from_reader(file)?) +} +/// Resolve the symbolic specs in `specs` and flatten them into a runtime +/// [`FormFactorMap`]. The inverse second half of [`parse_specs`]. +fn finalize_specs(mut specs: BTreeMap) -> Result { resolve_symbolic_specs(&mut specs, SpecRole::FormFactor)?; resolve_symbolic_specs(&mut specs, SpecRole::ExcludedSolvent)?; @@ -425,7 +439,12 @@ pub fn read_formfactors(path: &Path) -> Result { }) .collect::>()?; - let formfactor = FormFactorMap::from_entries_with_excluded(entries); + Ok(FormFactorMap::from_entries_with_excluded(entries)) +} + +/// Read YAML file with form factors for different atom types (atomic or coarse grained) +pub fn read_formfactors(path: &Path) -> Result { + let formfactor = finalize_specs(parse_specs(path)?)?; log::info!( "Read form-factor definitions from {} for {} species", path.display(), @@ -434,6 +453,48 @@ pub fn read_formfactors(path: &Path) -> Result { Ok(formfactor) } +/// Read and merge several YAML form-factor files into one table. +/// +/// Files are parsed in the given order and their species merged by name. +/// When the same species appears in more than one file, the **later** file +/// wins (last definition supplied), and a warning is logged naming the +/// species and the file that overrode it. Merging happens before symbolic +/// resolution, so an `!Alias` / `!United` / `!Linear` spec in one file may +/// reference a definition contributed by an earlier file. +pub fn read_formfactors_merged(paths: &[PathBuf]) -> Result { + let [first, rest @ ..] = paths else { + return Err(Error::validation("no form-factor files supplied")); + }; + + let mut merged = parse_specs(first)?; + for path in rest { + for (name, spec) in parse_specs(path)? { + use std::collections::btree_map::Entry; + match merged.entry(name) { + Entry::Occupied(mut entry) => { + log::warn!( + "Duplicate species '{}' redefined by {}; using the latest definition", + entry.key(), + path.display(), + ); + entry.insert(spec); + } + Entry::Vacant(entry) => { + entry.insert(spec); + } + } + } + } + + let formfactor = finalize_specs(merged)?; + log::info!( + "Merged form-factor definitions from {} file(s) into {} species", + paths.len(), + formfactor.len(), + ); + Ok(formfactor) +} + /// Convert all formfactors to a TABLE and Write to YAML file /// /// The q-interval is hard-coded to [0.0, 0.5] with 100 points. @@ -751,6 +812,46 @@ ALA: assert_relative_eq!(excluded.eval(0.2).unwrap(), 3.5); } + #[test] + fn read_formfactors_merged_latest_wins_and_cross_file_alias() { + // `base` defines C; `override_` redefines C (latest wins) and adds + // an alias that resolves against the base file's C — exercising the + // pre-resolution merge. + let base = r" +C: + formfactor: !Constant 6.0 +N: + formfactor: !Constant 7.0 +"; + let override_ = r" +C: + formfactor: !Constant 60.0 +CA: + formfactor: !Alias N +"; + let dir = std::env::temp_dir(); + let base_path = dir.join("pripps_merge_base.yaml"); + let over_path = dir.join("pripps_merge_override.yaml"); + std::fs::write(&base_path, base).unwrap(); + std::fs::write(&over_path, override_).unwrap(); + + let ff = read_formfactors_merged(&[base_path, over_path]).unwrap(); + + // Union of species: C, N (from base) + CA (from override). + assert_eq!(ff.len(), 3); + let eval = |name: &str| { + ff.form_factor_by_name(name) + .unwrap_or_else(|| panic!("missing {name}")) + .eval(0.0) + .unwrap() + }; + // Duplicate C: the later file's value wins. + assert_relative_eq!(eval("C"), 60.0); + assert_relative_eq!(eval("N"), 7.0); + // Alias in the override resolves against base's N. + assert_relative_eq!(eval("CA"), 7.0); + } + #[test] fn read_formfactors_resolves_symbolic_excluded_solvent() { let yaml = r" diff --git a/src/lib.rs b/src/lib.rs index 0d7783b..2ff0fbf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -55,7 +55,9 @@ pub use solvent::{ExcludedVolume, Hydration, SolventConfig}; // Crate-internal: the form-factor table and atom-kind types are part of // no public signature, so they stay off the public API surface. pub(crate) use atomkind::{AtomKind, FindByName}; -pub(crate) use formfactor::{FormFactorMap, read_formfactors, write_formfactors}; +pub(crate) use formfactor::{ + FormFactorMap, read_formfactors, read_formfactors_merged, write_formfactors, +}; /// Crate-internal short alias. Not re-exported; public signatures use /// `[f64; 3]` instead so the API doesn't leak `nalgebra` types. @@ -198,10 +200,21 @@ pub struct TrajectoryOptions<'a> { } impl Pripps { - /// Construct from a YAML form-factor definition file. + /// Construct from a single YAML form-factor definition file. pub fn from_yaml(path: &Path) -> Result { - let formfactors = read_formfactors(path)?; + Self::from_formfactors(read_formfactors(path)?) + } + + /// Construct from one or more YAML form-factor definition files, + /// merged in order. When a species is defined in more than one file + /// the later file wins and a warning is logged (see + /// [`read_formfactors_merged`]). Symbolic specs (`!Alias` / `!United` + /// / `!Linear`) may reference definitions from any of the files. + pub fn from_yaml_files(paths: &[std::path::PathBuf]) -> Result { + Self::from_formfactors(read_formfactors_merged(paths)?) + } + fn from_formfactors(formfactors: FormFactorMap) -> Result { if std::env::var("FF_TABLE").is_ok() { write_formfactors(&formfactors, Path::new("tabulated_formfactors.yaml"))?; } @@ -897,7 +910,9 @@ mod tests { /// rescales it back to a full average (not frame 0 × 1/N_frames). #[test] fn trajectory_stride_and_weights() { - let pripps = Pripps::from_yaml(Path::new("assets/gaussian_atoms.yaml")).unwrap(); + // 4lzt.xyz is a residue-bead structure, so it needs the CG + // single-bead form factors from poly_amino_acid.yaml. + let pripps = Pripps::from_yaml(Path::new("assets/poly_amino_acid.yaml")).unwrap(); let template = std::path::PathBuf::from("tests/lysozyme/4lzt.xyz"); let xtc = std::env::temp_dir().join("pripps_stride_test.xtc"); write_two_frame_xtc(&xtc); From 21580ec28c0f7a73a7fdceaac7e6ce44f874233d Mon Sep 17 00:00:00 2001 From: mlund Date: Fri, 19 Jun 2026 09:44:17 +0200 Subject: [PATCH 5/7] Address PR #14 review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename gaussian_atoms.yaml -> pepsi_atoms.yaml and document provenance: atomic five-Gaussians are Waasmaier & Kirfel (1995); united-atom groups and K=4π solvent convention are PepsiSAXS (Grudinin et al. 2017). Update all references (tests, script, README) and link DOIs inline. - Trim the README "Merging multiple files" section per review. - Use tempfile::tempdir() in the merge test to avoid fixed temp-file names colliding under parallel test runs. --- README.md | 27 ++++++------------- .../{gaussian_atoms.yaml => pepsi_atoms.yaml} | 9 ++++++- pripps-py/tests/test_smoke.py | 2 +- scripts/sasbdb_bench.py | 2 +- src/formfactor.rs | 6 ++--- src/lib.rs | 2 +- 6 files changed, 22 insertions(+), 26 deletions(-) rename assets/{gaussian_atoms.yaml => pepsi_atoms.yaml} (95%) diff --git a/README.md b/README.md index 38fa356..c140b1f 100644 --- a/README.md +++ b/README.md @@ -236,22 +236,11 @@ must define it when `--excluded per-kind` is selected. #### Merging multiple files -Several form-factor files can be supplied and are merged in the order -given — repeat `--atomfile` (or pass several paths) on the CLI, or pass a -list to the Python constructor: - -```bash -pripps --atomfile assets/gaussian_atoms.yaml --atomfile my_overrides.yaml ... -``` - -```python -pr = pripps.Pripps(["assets/gaussian_atoms.yaml", "my_overrides.yaml"]) -``` - -This lets you compose, e.g., a coarse-grained residue file with an atomic -file. **Duplicate policy:** if the same species is defined in more than one -file, the **last file listed wins** and a warning is logged naming the -overridden species. Merging happens before symbolic resolution, so an +Several form-factor files can be supplied and are merged in order — repeat +`--atomfile` on the CLI, or pass a list to the Python constructor — letting +you compose, e.g., a coarse-grained residue file with an atomic one. If the +same species appears in more than one file the **last file listed wins** and +a warning is logged. Merging happens before symbolic resolution, so an `!Alias` / `!United` / `!Linear` entry in one file may reference a species defined in another. @@ -260,8 +249,8 @@ defined in another. | File | Description | Best for | |------|-------------|----------| | `assets/foxs_formfactors.yaml` | Cromer-Mann (Int. Tables Vol. C) atomic form factors, FoXS-compatible volumes and SASA radii | Atomic PDBs with `--excluded foxs` | -| `assets/gaussian_atoms.yaml` | Atomic and united-atom Gaussians (CH/CH₂/NH/…), PepsiSAXS-style (K = 4π) | Atomic PDBs with PepsiSAXS-style solvent models | -| `assets/stovgaard2010.yaml` | Per-residue form-factor centroids (one-body model, q ∈ [0, 0.75] Å⁻¹) from Stovgaard et al. 2010, | Residue / Cα single-bead models | +| `assets/pepsi_atoms.yaml` | Atomic [Waasmaier-Kirfel](https://doi.org/10.1107/S0108767394013292) five-Gaussians + [PepsiSAXS](https://doi.org/10.1107/S2059798317005745) united-atom groups (CH/CH₂/NH/…), K = 4π | Atomic PDBs with PepsiSAXS-style solvent models | +| `assets/stovgaard2010.yaml` | Per-residue form-factor centroids (one-body model, q ∈ [0, 0.75] Å⁻¹) from [Stovgaard et al. 2010](https://doi.org/10.1186/1471-2105-11-429) | Residue / Cα single-bead models | | `assets/poly_amino_acid.yaml` | Amino-acid bead polynomials with a separate `excluded_solvent` polynomial per bead | Cα / single-bead with `--excluded per-kind` | | `assets/poly_amino_acid_combined.yaml` | Amino-acid bead polynomials with the excluded solvent already folded into the form factor | Cα / single-bead with `--excluded disabled` | | `assets/poly_martini.yaml` | Martini-3 bead polynomials with a separate `excluded_solvent` polynomial per bead | Martini-3 CG with `--excluded per-kind` | @@ -275,7 +264,7 @@ their respective tools. **Don't mix files with non-matching `--excluded` flags.** Each YAML + excluded-volume convention is a self-consistent bundle: `foxs_formfactors.yaml` was calibrated jointly with FoXS's -`K = 16π` Gaussian, while `gaussian_atoms.yaml` was calibrated with +`K = 16π` Gaussian, while `pepsi_atoms.yaml` was calibrated with PepsiSAXS's `K = 4π`. Swapping the denominator gives atomic and excluded contributions that were never fit together and yields subtly biased intensities. diff --git a/assets/gaussian_atoms.yaml b/assets/pepsi_atoms.yaml similarity index 95% rename from assets/gaussian_atoms.yaml rename to assets/pepsi_atoms.yaml index 1da56a2..bbf222a 100644 --- a/assets/gaussian_atoms.yaml +++ b/assets/pepsi_atoms.yaml @@ -1,4 +1,11 @@ -# Atomic / united-atom form factors for PRIPPS (PepsiSAXS-style K = 4π). +# Atomic / united-atom form factors for PRIPPS (PepsiSAXS-style, K = 4π). +# +# Atomic five-Gaussian scattering factors are the Waasmaier & Kirfel (1995) +# parametrization, valid for sinθ/λ ∈ [0, 6] Å⁻¹: +# https://doi.org/10.1107/S0108767394013292 +# combined with PepsiSAXS united-atom groups (CH/CH₂/NH/…) and the K = 4π +# solvent convention, Grudinin et al. (2017): +# https://doi.org/10.1107/S2059798317005745 # # Atomic-resolution structures only. Residue single-bead form factors live # in assets/poly_amino_acid.yaml — pass several --atomfile paths to merge diff --git a/pripps-py/tests/test_smoke.py b/pripps-py/tests/test_smoke.py index 39f1ac2..8b207c1 100644 --- a/pripps-py/tests/test_smoke.py +++ b/pripps-py/tests/test_smoke.py @@ -132,7 +132,7 @@ class TestPepsiFit(unittest.TestCase): def setUpClass(cls): cls.exp = load_lyzexp() q_max = float(cls.exp[-1, 0]) - pr = pripps.Pripps("../assets/gaussian_atoms.yaml") + pr = pripps.Pripps("../assets/pepsi_atoms.yaml") cls.model = pr.multipole_model( "../tests/lysozyme/4lzt.pdb", qmin=0.0, diff --git a/scripts/sasbdb_bench.py b/scripts/sasbdb_bench.py index 569529f..fe2c956 100755 --- a/scripts/sasbdb_bench.py +++ b/scripts/sasbdb_bench.py @@ -62,7 +62,7 @@ BENCH = REPO / "sasbdb_bench" CACHE = BENCH / "cache" FF_YAML_DEFAULT = REPO / "assets" / "foxs_formfactors.yaml" -# Override via --atomfile if you want gaussian_atoms.yaml for the +# Override via --atomfile if you want pepsi_atoms.yaml for the # PepsiSAXS-style K=4π convention. SASBDB_API = "https://www.sasbdb.org/rest-api/entry/summary/{}/" diff --git a/src/formfactor.rs b/src/formfactor.rs index 4a85305..6020549 100644 --- a/src/formfactor.rs +++ b/src/formfactor.rs @@ -829,9 +829,9 @@ C: CA: formfactor: !Alias N "; - let dir = std::env::temp_dir(); - let base_path = dir.join("pripps_merge_base.yaml"); - let over_path = dir.join("pripps_merge_override.yaml"); + let dir = tempfile::tempdir().unwrap(); + let base_path = dir.path().join("base.yaml"); + let over_path = dir.path().join("override.yaml"); std::fs::write(&base_path, base).unwrap(); std::fs::write(&over_path, override_).unwrap(); diff --git a/src/lib.rs b/src/lib.rs index 2ff0fbf..061741b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -837,7 +837,7 @@ mod tests { // produce a physically reasonable profile. Not as good as // the FoXS pipeline (different coefficients, different // denominator) but the shape should be plausible. - let pripps = Pripps::from_yaml(std::path::Path::new("assets/gaussian_atoms.yaml")).unwrap(); + let pripps = Pripps::from_yaml(std::path::Path::new("assets/pepsi_atoms.yaml")).unwrap(); let pdb = std::path::PathBuf::from("tests/lysozyme/4lzt.pdb"); let solvent = SolventConfig { excluded: ExcludedVolume::Fraser { volume_scale: 1.0 }, From 851ad9bb4a565287a9fd55d5011dbd094d380371 Mon Sep 17 00:00:00 2001 From: mlund Date: Fri, 19 Jun 2026 10:00:19 +0200 Subject: [PATCH 6/7] Clarify gaussian ff in readme --- README.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index c140b1f..e6c1d9e 100644 --- a/README.md +++ b/README.md @@ -256,10 +256,12 @@ defined in another. | `assets/poly_martini.yaml` | Martini-3 bead polynomials with a separate `excluded_solvent` polynomial per bead | Martini-3 CG with `--excluded per-kind` | | `assets/poly_martini_combined.yaml` | Martini-3 bead polynomials with the excluded solvent already folded into the form factor | Martini-3 CG with `--excluded disabled` | -Gaussian files use b values in the standard s² = (q/4π)² convention -(International Tables / Cromer-Mann). The atomic coefficients differ -(Cromer-Mann vs Stovgaard-2010) and the volume/radius tables match -their respective tools. +Files using `!Gaussian` form factors (`foxs_formfactors.yaml`, +`pepsi_atoms.yaml`) give b values in the standard s² = (q/4π)² convention +(International Tables / Cromer-Mann), so published exponents can be used +directly. Their atomic coefficients come from different parametrizations +(Cromer-Mann for FoXS vs Waasmaier-Kirfel for Pepsi) and the volume/radius +tables match their respective tools. **Don't mix files with non-matching `--excluded` flags.** Each YAML + excluded-volume convention is a self-consistent bundle: From 69bdde45c62f80b906b089256faa3fe115867c89 Mon Sep 17 00:00:00 2001 From: Mikael Lund Date: Mon, 22 Jun 2026 10:56:18 +0200 Subject: [PATCH 7/7] Number of molecules input; S_eff(q) output (#15) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add --num-molecules: effective structure factor S_eff(q) For a periodic box of N identical molecules (the `direct` scheme, optionally XTC-averaged), `--num-molecules N` outputs the apparent structure factor and the at-concentration single-molecule form factor: S_eff(q) = |Σ_m F_m(q)|² / Σ_m |F_m(q)|² (exact, no approximation) P(q) = (1/N) · Σ_m |F_m(q)|² The box is split into N contiguous equal-size molecules and each molecule's amplitude F_m is accumulated alongside the self term Σ_m|F_m|² in DirectTransform::vacuum_intensities. Both use the box-commensurate q-vectors, so they are PBC-invariant (no unwrapping). S_eff and P are derived from the self term in Pripps::intensity and emitted as the CSV columns `p,s_eff` (and the Python .p/.s_eff getters). P(q) is the at-concentration form factor — for a flexible molecule it shifts with crowding, unlike the dilute P an experiment divides by. Guards (fail up front, before any trajectory pass): direct scheme only, vacuum only, box atom count a multiple of N, and every M-atom block must share the first block's atom-kind sequence (rejects a count that splits molecules or doesn't match the box composition). --num-molecules conflicts with --fit and runs on the CPU even with the gpu feature (the GPU kernel emits only the total |Σf|²); both are noted. XTC averaging assumes a fixed box. * Run --num-molecules on the GPU; tidy and test Extend the WGSL direct kernel to emit the per-molecule self term Σ_m|F_m|² alongside the box intensity (one workgroup per q-vector, threads stride over whole molecules), so the effective structure factor runs on the GPU instead of falling back to the CPU. Output is two f32 per q-vector; the whole-box path (no --num-molecules) is unchanged. Route DirectGpu + --num-molecules to the GPU, with the existing CPU fallback when no adapter is available. Cleanups: route average_duplicates and the 5-column averager through the shared average_by_qnorm (removing the triplicated q-binning/sort), and add an OUT_BYTES_PER_Q constant for the GPU output stride. Rename the internal Intensity carrier monomer_self -> self_term to match the local accumulators and drop the stale "monomer" naming. Tests: GPU-vs-CPU parity for the self term, and a trajectory test pinning S_eff as the ratio of frame-averaged box intensity to frame-averaged self term (ratio of averages, not the average of per-frame S_eff). * Address PR #15 review: guard zero consumed-weight-sum trajectory_average divides accumulated curves by the consumed weight sum. read_weights validates the full set, but striding can select a zero-sum subset of signed weights, which would turn every average into NaN; reject it explicitly instead. The companion review note about an empty box (0 atoms) yielding NaN S_eff is already covered upstream — Structure::from_file rejects atomless inputs before the --num-molecules guard runs — so no extra check is needed there; add a regression test and a clarifying comment. * README: derive the effective structure factor theory Replace the terse --num-molecules block with a proper derivation (GitHub LaTeX): per-molecule amplitude, the self/cross split of the box intensity, the exact P(q) and S_eff(q) definitions, and the decoupling relation S_eff = 1 + beta(S-1) with the beta-node inversion caveat. Reference the primary source for the decoupling approximation (Kotlarchyk & Chen, J. Chem. Phys. 79, 2461 (1983)). * README: fix inline math rendering on GitHub GitHub's markdown pairs the underscores in adjacent inline $...$ spans as emphasis, stripping them and breaking the math (e.g. R_m, S_eff). Switch inline math to the backtick-protected $`...`$ form, whose contents are immune to markdown; display blocks are unchanged. --- README.md | 101 ++++++++++++++ benches/intensity.rs | 12 +- pripps-py/src/lib.rs | 22 +++ src/cli.rs | 48 ++++--- src/debye.rs | 1 + src/explicit.rs | 117 ++++++++++------ src/gpu.rs | 187 ++++++++++++++++++++----- src/lib.rs | 324 +++++++++++++++++++++++++++++++++++++++++-- src/multipole/mod.rs | 2 + src/trajectory.rs | 30 +++- 10 files changed, 738 insertions(+), 106 deletions(-) diff --git a/README.md b/README.md index e6c1d9e..e68f505 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,107 @@ avg = pr.intensity( print(avg.q, avg.total) ``` +### Effective structure factor (`--num-molecules`) + +For a periodic box of $`N`$ identical molecules — a concentrated solution +simulated with the `direct` scheme, optionally averaged over an XTC trajectory — +`--num-molecules N` (CLI) or `num_molecules=N` (Python) separates the +single-molecule **form factor** $`P(q)`$ from the inter-molecular **structure +factor**. + +Each molecule $`m`$ scatters with amplitude + +```math +F_m(\mathbf{q}) = \sum_{i \in m} f_i(q)\, e^{\,i\,\mathbf{q}\cdot\mathbf{r}_i}, +``` + +where $`f_i(q)`$ is the atomic form factor and $`\mathbf{r}_i`$ the atom position. +The box intensity is the squared modulus of the total amplitude, which splits +into intra-molecular (self) and inter-molecular (cross) terms: + +```math +I(\mathbf{q}) = \left|\sum_{m=1}^{N} F_m(\mathbf{q})\right|^2 + = \underbrace{\sum_{m} \left|F_m\right|^2}_{\text{self}} + + \underbrace{\sum_{m \neq m'} F_m F_{m'}^{*}}_{\text{cross}} . +``` + +Writing $`\langle\,\cdot\,\rangle`$ for the average over scattering-vector +orientations at fixed $`q = |\mathbf{q}|`$ (and over trajectory frames), the form +factor and the **effective** (apparent) structure factor are + +```math +P(q) = \frac{1}{N}\left\langle \sum_{m} \left|F_m(\mathbf{q})\right|^2 \right\rangle + = \bigl\langle \left|F(\mathbf{q})\right|^2 \bigr\rangle_{\Omega}, +\qquad +S_\mathrm{eff}(q) = \frac{\langle I(\mathbf{q})\rangle}{N\,P(q)} + = \frac{\bigl\langle \left|\sum_m F_m\right|^2 \bigr\rangle} + {\bigl\langle \sum_m \left|F_m\right|^2 \bigr\rangle} . +``` + +pripps emits exactly these two columns, `p` and `s_eff`. $`S_\mathrm{eff}(q)`$ is +the quantity a scattering experiment measures — the box intensity normalised by +$`N`$ independent molecules — and is obtained directly from the self/cross split, +so it is **exact**: no assumption about particle shape or correlations is made. + +**Relation to the true structure factor.** The thermodynamic centre-to-centre +structure factor, the Fourier transform of the centre–centre pair correlation +$`g(r)`$, is + +```math +S(q) = 1 + \frac{1}{N}\left\langle \sum_{m \neq m'} + e^{\,i\,\mathbf{q}\cdot(\mathbf{R}_m - \mathbf{R}_{m'})} \right\rangle, +``` + +with $`\mathbf{R}_m`$ the molecular centre. For non-spherical molecules +$`S_\mathrm{eff} \neq S`$. Under the *decoupling approximation* — orientation +uncorrelated with position — they are related by + +```math +S_\mathrm{eff}(q) = 1 + \beta(q)\,\bigl[S(q) - 1\bigr], +\qquad +\beta(q) = \frac{\bigl|\langle F(\mathbf{q})\rangle_{\Omega}\bigr|^2} + {\bigl\langle |F(\mathbf{q})|^2 \bigr\rangle_{\Omega}}, +``` + +where the decoupling parameter obeys $`0 \le \beta(q) \le 1`$, with $`\beta = 1`$ for +centrosymmetric (e.g. spherical) particles, giving $`S_\mathrm{eff} = S`$. The true +$`S(q)`$ follows by inversion, + +```math +S(q) = 1 + \frac{S_\mathrm{eff}(q) - 1}{\beta(q)}, +``` + +which is ill-conditioned near the nodes of $`\beta(q)`$; pripps therefore reports +$`P`$ and $`S_\mathrm{eff}`$ and leaves $`\beta(q)`$ and the (judgement-dependent) +inversion to the user. See Kotlarchyk & Chen, *J. Chem. Phys.* **79**, 2461 +(1983), [doi:10.1063/1.446055](https://doi.org/10.1063/1.446055). + +**Notes.** +- The amplitudes use the box-commensurate vectors $`\mathbf{q} = 2\pi p\,(h,k,l)/L`$, + for which $`e^{\,i\,\mathbf{q}\cdot\mathbf{r}}`$ is unchanged under + $`\mathbf{r}\to\mathbf{r}+\mathbf{L}`$. Both $`I`$ and the self term are thus + PBC-invariant and need no coordinate unwrapping; the orientational average is + built from the sampled $`(h,k,l)`$ directions, and the trajectory average is a + ratio of frame-averaged numerator and denominator. +- $`P(q)`$ is the **at-concentration** form factor — the molecules as they actually + are in the crowded box. For a flexible molecule it depends on concentration, + unlike the dilute-limit $`P`$ an experiment divides out, so comparing the two + reveals conformational change. It is on an absolute scale, + $`P(0) = \bigl(\sum_i f_i\bigr)^2`$, not normalised to $`P(0)=1`$. + +```sh +pripps -a atoms.yaml -i box.xyz -x traj.xtc --num-molecules 100 direct -p 100 +``` + +```python +out = pr.intensity("box.xyz", model="direct", box_sides=200.0, num_molecules=100) +print(out.s_eff, out.p) +``` + +The box atom count must be a multiple of $`N`$ (contiguous, identical molecules), +and the run must be vacuum. The `gpu` feature accelerates this path too; XTC +averaging assumes a fixed box (constant volume). + ### Solvent models | Parameter | Values | Description | diff --git a/benches/intensity.rs b/benches/intensity.rs index 74f902d..836e3c5 100644 --- a/benches/intensity.rs +++ b/benches/intensity.rs @@ -164,7 +164,11 @@ fn bench_intensity(c: &mut Criterion) { // unused. Without this the optimizer could realize // that each iteration is idempotent and skip most of // the work. - black_box(pripps.intensity(scheme, pdb, None, None, None).unwrap()); + black_box( + pripps + .intensity(scheme, pdb, None, None, None, None) + .unwrap(), + ); }); }); @@ -178,7 +182,11 @@ fn bench_intensity(c: &mut Criterion) { // matching the default CLI behaviour. l_max: None, }; - black_box(pripps.intensity(scheme, pdb, None, None, None).unwrap()); + black_box( + pripps + .intensity(scheme, pdb, None, None, None, None) + .unwrap(), + ); }); }); } diff --git a/pripps-py/src/lib.rs b/pripps-py/src/lib.rs index aecfda5..5449731 100644 --- a/pripps-py/src/lib.rs +++ b/pripps-py/src/lib.rs @@ -142,12 +142,15 @@ struct PyIntensity { atoms: Option>>, excluded: Option>>, hydration: Option>>, + p: Option>>, + s_eff: Option>>, n_points: usize, } impl PyIntensity { fn new(py: Python<'_>, inner: Intensity) -> Self { let arr = |v: Vec| v.into_pyarray(py).unbind(); + let opt = |v: Option>| v.map(arr); let n_points = inner.q.len(); let (atoms, excluded, hydration) = match inner.contributions { Some(c) => ( @@ -163,6 +166,8 @@ impl PyIntensity { atoms, excluded, hydration, + p: opt(inner.form_factor), + s_eff: opt(inner.effective_structure_factor), n_points, } } @@ -200,6 +205,20 @@ impl PyIntensity { self.hydration.as_ref().map(|a| a.clone_ref(py)) } + /// At-concentration single-molecule form factor P(q), or `None` unless + /// `num_molecules` was given. + #[getter] + fn p(&self, py: Python<'_>) -> Option>> { + self.p.as_ref().map(|a| a.clone_ref(py)) + } + + /// Effective structure factor S_eff(q) = I/(N·P), or `None` unless + /// `num_molecules` was given. + #[getter] + fn s_eff(&self, py: Python<'_>) -> Option>> { + self.s_eff.as_ref().map(|a| a.clone_ref(py)) + } + fn __repr__(&self) -> String { format!( "Intensity(points={}, contributions={})", @@ -463,6 +482,7 @@ impl PyPripps { bulk_electron_density = 0.334, volume_scale = 1.0, contrast_density = 0.03, + num_molecules = None, ))] fn intensity( &self, @@ -487,6 +507,7 @@ impl PyPripps { bulk_electron_density: f64, volume_scale: f64, contrast_density: f64, + num_molecules: Option, ) -> PyResult { let scheme = match scheme_from(model.as_ref())? { PyScheme::Multipole => IntensityScheme::Multipole { @@ -533,6 +554,7 @@ impl PyPripps { xtc_opts, box_sides, Some(&cfg), + num_molecules, ) .map(|inner| PyIntensity::new(py, inner)) .map_err(into_py_err) diff --git a/src/cli.rs b/src/cli.rs index eb5391b..3fa62a9 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -234,6 +234,10 @@ pub struct Args { #[clap(short = 'f', long)] fit: Option, + /// Number of identical molecules in the box → add S_eff(q) (direct, vacuum) + #[clap(long, conflicts_with = "fit")] + num_molecules: Option, + /// Optional .xtc trajectory file with structures to average over #[clap(short = 'x', long)] xtcfile: Option, @@ -345,6 +349,7 @@ pub fn do_main() -> Result<()> { xtc_opts, box_override, solvent.as_ref(), + args.num_molecules, )?; write_csv(&args.output, &intensity)?; } @@ -418,29 +423,30 @@ fn write_csv(path: &Path, curve: &Intensity) -> Result<()> { let mut out = BufWriter::new(file); log::info!("Writing {}", path.display()); - match &curve.contributions { - None => { - writeln!(out, "q,total")?; - for (q, total) in curve.q.iter().zip(&curve.total) { - writeln!(out, "{q:.6},{total:.6}")?; - } + // Header and rows gate the same optional blocks (contributions, then the + // `p,s_eff` structure-factor pair) so their column sets can't diverge. + let mut header = String::from("q,total"); + if curve.contributions.is_some() { + header.push_str(",atoms,excluded,hydration"); + } + if curve.effective_structure_factor.is_some() { + header.push_str(",p,s_eff"); + } + writeln!(out, "{header}")?; + + for i in 0..curve.q.len() { + write!(out, "{:.6},{:.6}", curve.q[i], curve.total[i])?; + if let Some(c) = &curve.contributions { + write!( + out, + ",{:.6},{:.6},{:.6}", + c.atoms[i], c.excluded[i], c.hydration[i] + )?; } - Some(c) => { - writeln!(out, "q,total,atoms,excluded,hydration")?; - for ((((q, total), atoms), excluded), hydration) in curve - .q - .iter() - .zip(&curve.total) - .zip(&c.atoms) - .zip(&c.excluded) - .zip(&c.hydration) - { - writeln!( - out, - "{q:.6},{total:.6},{atoms:.6},{excluded:.6},{hydration:.6}" - )?; - } + if let (Some(p), Some(s)) = (&curve.form_factor, &curve.effective_structure_factor) { + write!(out, ",{:.6},{:.6}", p[i], s[i])?; } + writeln!(out)?; } Ok(()) } diff --git a/src/debye.rs b/src/debye.rs index 0cba383..06c33bf 100644 --- a/src/debye.rs +++ b/src/debye.rs @@ -163,6 +163,7 @@ impl DebyeModel { excluded, hydration, }), + ..Default::default() } } diff --git a/src/explicit.rs b/src/explicit.rs index 4203508..8737263 100644 --- a/src/explicit.rs +++ b/src/explicit.rs @@ -22,7 +22,6 @@ use crate::{ use itertools::Itertools; use num_complex::Complex; use rayon::prelude::*; -use std::ops::Mul; use std::{cmp::Ordering, f64::consts::PI}; /// Reciprocal-space sampling directions (h, k, l). @@ -74,6 +73,10 @@ pub(crate) struct DirectTransform { formfactors: FormFactorMap, /// Optional solvent correction solvent: Option, + /// When set, the box is split into this many contiguous, equal-size + /// molecules so the per-molecule self term `Σ_m|F_m|²` is accumulated + /// alongside `I_box = |Σ_m F_m|²` (see [`Intensity::self_term`]). + num_molecules: Option, } /// Complex amplitude A(q) = Σ f_i · exp(i·q·r_i) for a set of @@ -127,25 +130,51 @@ impl DirectTransform { formfactors, directions, solvent, + num_molecules: None, } } + /// Split the box into `n` contiguous equal-size molecules so each + /// calculation also emits the self term `Σ_m|F_m|²` in + /// [`Intensity::self_term`] (for the effective structure factor). + pub(crate) fn with_num_molecules(mut self, n: usize) -> Self { + self.num_molecules = Some(n); + self + } + fn vacuum_intensities(&self, structure: &Structure, box_sides: &Vector3) -> Result { - let table = self + // One path for both modes: `n = 1` is the whole box, so the per-molecule + // sum collapses to the plain box amplitude and the extra self term is + // dropped below. Commensurate q-vectors make `|Σ_m F_m|²` and `Σ_m|F_m|²` + // PBC-invariant, so the molecule positions never need unwrapping. + let n = self.num_molecules.unwrap_or(1); + let m = (structure.pos.len() / n).max(1); + let rows: Vec<(f64, [f64; 2])> = self .directions .par_iter() .map(|dir| dir.component_div(box_sides)) .map(|q| { let ff = eval_atom_form_factors(q.norm(), &structure.ids, &self.formfactors)?; - let intensity = complex_amplitude(&q, &structure.pos, &ff).norm_sqr(); - Ok((q.norm(), intensity)) + let mut a = Complex::new(0.0, 0.0); + let mut self_term = 0.0; + for (pos, ff) in structure.pos.chunks(m).zip(ff.chunks(m)) { + let f_m = complex_amplitude(&q, pos, ff); + a += f_m; + self_term += f_m.norm_sqr(); + } + Ok((q.norm(), [a.norm_sqr(), self_term])) }) - .collect::>>()?; + .collect::>>()?; + let grouped = average_by_qnorm(rows); - let (q, total): (Vec, Vec) = average_duplicates(table).into_iter().unzip(); Ok(Intensity { - q, - total, + q: grouped.iter().map(|g| g.0).collect(), + total: grouped.iter().map(|g| g.1[0]).collect(), + // The self term is only meaningful (and only output) when the box + // was split into molecules; for the whole box it equals `total`. + self_term: self + .num_molecules + .map(|_| grouped.iter().map(|g| g.1[1]).collect()), ..Default::default() }) } @@ -279,52 +308,61 @@ fn eval_atom_form_factors( Ok(atom_ids.iter().map(|&id| per_kind[id]).collect()) } -/// Merge duplicate q-values by averaging, for 2-column data. -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::() / pair.len() as f64; - (pair.first().unwrap().0, mean_y) - }; - - let mut table: Vec<_> = data - .into_iter() - .into_group_map_by(|(x, _)| round(*x)) - .values() - .map(|pair| average(pair)) - .collect(); - table.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); - table +/// Bin a q-magnitude for grouping duplicate q-values (q-vectors of the same +/// length along different Miller directions). The single rounding tolerance +/// used by every `average_*` helper here. +fn q_bin(q: f64) -> i64 { + (q * 1e6).round() as i64 } -/// Merge duplicate q-values by averaging, for 5-column data. -fn average_duplicates_5col(data: Vec<(f64, f64, f64, f64, f64)>) -> Vec<(f64, f64, f64, f64, f64)> { - let round = |x: f64| x.mul(1e6).round() as i64; - data.into_iter() - .into_group_map_by(|(q, _, _, _, _)| round(*q)) +/// Merge duplicate q-values by averaging a fixed-width row of accumulated +/// columns, sorted by q. Used for the per-molecule (I_box, self) pair on both +/// the CPU and GPU direct paths. +pub(crate) fn average_by_qnorm(rows: Vec<(f64, [f64; C])>) -> Vec<(f64, [f64; C])> { + let mut out: Vec<(f64, [f64; C])> = rows + .into_iter() + .into_group_map_by(|(qn, _)| q_bin(*qn)) .into_values() .map(|group| { - let n = group.len() as f64; - let q = group[0].0; - let t = group.iter().map(|g| g.1).sum::() / n; - let a = group.iter().map(|g| g.2).sum::() / n; - let e = group.iter().map(|g| g.3).sum::() / n; - let h = group.iter().map(|g| g.4).sum::() / n; - (q, t, a, e, h) + let len = group.len() as f64; + let mut acc = [0.0; C]; + for (_, v) in &group { + for (a, &x) in acc.iter_mut().zip(v) { + *a += x; + } + } + acc.iter_mut().for_each(|a| *a /= len); + (group[0].0, acc) }) + .collect(); + out.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + out +} + +/// Merge duplicate q-values by averaging, for 2-column data. Thin wrapper over +/// [`average_by_qnorm`] so the grouping/binning/sort lives in one place. +pub(crate) fn average_duplicates(data: Vec<(f64, f64)>) -> Vec<(f64, f64)> { + average_by_qnorm(data.into_iter().map(|(q, y)| (q, [y])).collect()) + .into_iter() + .map(|(q, [y])| (q, y)) .collect() } fn intensity_from_5col(table: Vec<(f64, f64, f64, f64, f64)>) -> Intensity { - let mut grouped = average_duplicates_5col(table); - grouped.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal)); + // Average the four data columns by q-magnitude (sorted) in one place. + let grouped = average_by_qnorm( + table + .into_iter() + .map(|(q, t, a, e, h)| (q, [t, a, e, h])) + .collect(), + ); let mut q = Vec::with_capacity(grouped.len()); let mut total = Vec::with_capacity(grouped.len()); let mut atoms = Vec::with_capacity(grouped.len()); let mut excluded = Vec::with_capacity(grouped.len()); let mut hydration = Vec::with_capacity(grouped.len()); - for (qi, t, a, e, h) in grouped { + for (qi, [t, a, e, h]) in grouped { q.push(qi); total.push(t); atoms.push(a); @@ -340,6 +378,7 @@ fn intensity_from_5col(table: Vec<(f64, f64, f64, f64, f64)>) -> Intensity { excluded, hydration, }), + ..Default::default() } } diff --git a/src/gpu.rs b/src/gpu.rs index c73aa4d..d3f3fb1 100644 --- a/src/gpu.rs +++ b/src/gpu.rs @@ -29,12 +29,16 @@ //! The win is **trajectory averaging**: the q-vectors and atom ids stay //! resident on the device across frames, and only the atom positions (and the //! per-frame box / form-factor table) are re-uploaded each frame. +//! +//! With [`with_num_molecules`](GpuDirectTransform::with_num_molecules) the +//! kernel additionally groups atoms into molecules and emits the per-molecule +//! self term `Σ_m|F_m|²` alongside `I_box`, for the effective structure factor. use std::sync::Mutex; use wgpu::util::DeviceExt; -use crate::explicit::{average_duplicates, directions}; +use crate::explicit::{average_by_qnorm, average_duplicates, directions}; use crate::formfactor::FormFactor; use crate::{Error, FormFactorMap, Intensity, IntensityCalculator, Result, Structure, Vector3}; @@ -46,8 +50,8 @@ struct Params { inv_box: vec4, // 1/box_sides; .xyz used n_atoms: u32, n_q: u32, - _pad0: u32, - _pad1: u32, + atoms_per_mol: u32, // M; 0 = whole box (self term unused) + _pad: u32, }; @group(0) @binding(0) var params: Params; @@ -55,10 +59,11 @@ struct Params { @group(0) @binding(2) var positions: array>; // n_atoms (.xyz used) @group(0) @binding(3) var atom_ids: array; // n_atoms @group(0) @binding(4) var ff_table: array; // n_kinds * n_q (kind-major) -@group(0) @binding(5) var out_iq: array; // n_q +@group(0) @binding(5) var out_iq: array; // 2 * n_q: [I_box, self] var sh_re: array; var sh_im: array; +var sh_self: array; @compute @workgroup_size(256) fn structure_factor( @@ -68,30 +73,66 @@ fn structure_factor( let qi = wg.x; let q = directions[qi].xyz * params.inv_box.xyz; - // Strided walk over atoms; each thread accumulates its own partial sum. + // total = Σ_m F_m (→ I_box = |total|²); self = Σ_m |F_m|². var re = 0.0; var im = 0.0; - var a = lid.x; - loop { - if (a >= params.n_atoms) { break; } - let r = positions[a].xyz; - let f = ff_table[atom_ids[a] * params.n_q + qi]; - let qr = dot(q, r); - re = re + f * cos(qr); - im = im + f * sin(qr); - a = a + 256u; + var self_term = 0.0; + + let m = params.atoms_per_mol; + if (m == 0u) { + // The original whole-box path (no --num-molecules); self term stays 0. + var a = lid.x; + loop { + if (a >= params.n_atoms) { break; } + let r = positions[a].xyz; + let f = ff_table[atom_ids[a] * params.n_q + qi]; + let qr = dot(q, r); + re = re + f * cos(qr); + im = im + f * sin(qr); + a = a + 256u; + } + } else { + // Per molecule: each thread strides over whole molecules (so each F_m + // is summed by a single thread); accumulate the box amplitude and the + // self term. With fewer than 256 molecules some threads idle, but the + // n_q workgroups keep the device busy. + let n_mol = params.n_atoms / m; + var mol = lid.x; + loop { + if (mol >= n_mol) { break; } + let start = mol * m; + var fre = 0.0; + var fim = 0.0; + var k = 0u; + loop { + if (k >= m) { break; } + let a = start + k; + let r = positions[a].xyz; + let f = ff_table[atom_ids[a] * params.n_q + qi]; + let qr = dot(q, r); + fre = fre + f * cos(qr); + fim = fim + f * sin(qr); + k = k + 1u; + } + re = re + fre; + im = im + fim; + self_term = self_term + fre * fre + fim * fim; + mol = mol + 256u; + } } sh_re[lid.x] = re; sh_im[lid.x] = im; + sh_self[lid.x] = self_term; workgroupBarrier(); - // Shared-memory tree reduction. + // Shared-memory tree reduction of all three partial sums. var stride = 128u; loop { if (stride == 0u) { break; } if (lid.x < stride) { sh_re[lid.x] = sh_re[lid.x] + sh_re[lid.x + stride]; sh_im[lid.x] = sh_im[lid.x] + sh_im[lid.x + stride]; + sh_self[lid.x] = sh_self[lid.x] + sh_self[lid.x + stride]; } workgroupBarrier(); stride = stride / 2u; @@ -100,23 +141,29 @@ fn structure_factor( if (lid.x == 0u) { let rr = sh_re[0]; let ii = sh_im[0]; - out_iq[qi] = rr * rr + ii * ii; + out_iq[2u * qi] = rr * rr + ii * ii; + out_iq[2u * qi + 1u] = sh_self[0]; } } "#; /// Per-frame uniform block. `#[repr(C)]` with explicit padding so the layout -/// matches the WGSL `Params` struct (vec4 + two u32 + two pad u32 = 32 bytes, -/// the size rounding required by the vec4's 16-byte alignment). +/// matches the WGSL `Params` struct (vec4 + four u32 = 32 bytes, the size +/// rounding required by the vec4's 16-byte alignment). #[repr(C)] #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] struct Params { inv_box: [f32; 4], n_atoms: u32, n_q: u32, - _pad: [u32; 2], + /// Atoms per molecule (`M`); 0 = whole box, the self term is then unused. + atoms_per_mol: u32, + _pad: u32, } +/// Bytes the kernel writes per q-vector: two `f32`, `[I_box, self term]`. +const OUT_BYTES_PER_Q: u64 = 2 * 4; + /// GPU device context plus the static, structure-independent resources /// (q-directions, pipeline). Built once in [`GpuDirectTransform::new`]. struct GpuContext { @@ -150,6 +197,9 @@ pub(crate) struct GpuDirectTransform { /// q-directions in host (f64) form, for per-frame `|q|` and box division. directions: Vec, ctx: GpuContext, + /// When set, split the box into this many contiguous molecules and emit the + /// per-molecule self term `Σ_m|F_m|²` for the effective structure factor. + num_molecules: Option, /// Lazily built on the first frame; the trait is `&self`, so interior /// mutability is required. The trajectory loop is serial, so contention /// is nil. @@ -225,9 +275,17 @@ impl GpuDirectTransform { dir_buf, n_q, }, + num_molecules: None, frame: Mutex::new(None), }) } + + /// Split the box into `n` contiguous equal-size molecules so each frame also + /// emits the self term `Σ_m|F_m|²` in [`Intensity::self_term`]. + pub(crate) fn with_num_molecules(mut self, n: usize) -> Self { + self.num_molecules = Some(n); + self + } } impl IntensityCalculator for GpuDirectTransform { @@ -286,6 +344,13 @@ impl IntensityCalculator for GpuDirectTransform { .map(|p| [p.x as f32, p.y as f32, p.z as f32, 0.0]) .collect(); + // Atoms per molecule (0 = whole box). Divisibility is enforced upstream + // by the `--num-molecules` guard in `Pripps::intensity`. + let atoms_per_mol = self + .num_molecules + .map(|n| structure.pos.len() / n) + .unwrap_or(0) as u32; + let params = Params { inv_box: [ (1.0 / box_sides.x) as f32, @@ -295,7 +360,8 @@ impl IntensityCalculator for GpuDirectTransform { ], n_atoms, n_q: ctx.n_q, - _pad: [0, 0], + atoms_per_mol, + _pad: 0, }; ctx.queue @@ -305,19 +371,23 @@ impl IntensityCalculator for GpuDirectTransform { ctx.queue .write_buffer(&frame.params_buf, 0, bytemuck::bytes_of(¶ms)); - let intensities = ctx.dispatch(frame)?; - - // Pair each |q| with its I(q) and average duplicate magnitudes, exactly - // as the CPU vacuum path does. - let table: Vec<(f64, f64)> = qmags + // The kernel writes two f32 per q-vector: [I_box, self term]. + let out = ctx.dispatch(frame)?; + let rows: Vec<(f64, [f64; 2])> = qmags .into_iter() - .zip(intensities.into_iter().map(f64::from)) + .enumerate() + .map(|(i, q)| (q, [f64::from(out[2 * i]), f64::from(out[2 * i + 1])])) .collect(); - let (q, total): (Vec, Vec) = average_duplicates(table).into_iter().unzip(); + // Average duplicate |q| magnitudes, exactly as the CPU vacuum path does. + let grouped = average_by_qnorm(rows); Ok(Intensity { - q, - total, + q: grouped.iter().map(|g| g.0).collect(), + total: grouped.iter().map(|g| g.1[0]).collect(), + // The self term is only meaningful when the box was split. + self_term: self + .num_molecules + .map(|_| grouped.iter().map(|g| g.1[1]).collect()), ..Default::default() }) } @@ -371,7 +441,8 @@ impl GpuContext { usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, mapped_at_creation: false, }); - let out_size = u64::from(self.n_q) * 4; + // Two f32 per q-vector: [I_box, self term]. + let out_size = u64::from(self.n_q) * OUT_BYTES_PER_Q; let out_buf = device.create_buffer(&wgpu::BufferDescriptor { label: Some("out-iq"), size: out_size, @@ -448,7 +519,7 @@ impl GpuContext { pass.set_bind_group(0, &frame.bind_group, &[]); pass.dispatch_workgroups(self.n_q, 1, 1); } - let out_size = u64::from(self.n_q) * 4; + let out_size = u64::from(self.n_q) * OUT_BYTES_PER_Q; encoder.copy_buffer_to_buffer(&frame.out_buf, 0, &frame.read_buf, 0, out_size); self.queue.submit([encoder.finish()]); @@ -545,4 +616,58 @@ mod tests { ); } } + + /// With `--num-molecules`, the GPU must also reproduce the CPU per-molecule + /// self term `Σ_m|F_m|²` (and hence S_eff/P). Compares the raw `total` and + /// `self_term`; the S_eff/P finalize is identical host arithmetic. + #[test] + fn gpu_matches_cpu_num_molecules() { + let dir = std::env::temp_dir(); + let yaml = dir.join("pripps_gpu_nmol.yaml"); + std::fs::write(&yaml, "X: { formfactor: !Constant 1.0 }\n").unwrap(); + // Three identical 2-atom molecules spread through the box. + let xyz = dir.join("pripps_gpu_nmol.xyz"); + std::fs::write( + &xyz, + "6\nb\nX 0 0 0\nX 0 0 4\nX 30 10 0\nX 30 10 4\nX 70 50 0\nX 70 50 4\n", + ) + .unwrap(); + let pripps = Pripps::from_yaml(&yaml).unwrap(); + + let gpu = match GpuDirectTransform::new(pripps.formfactors.clone(), 6) { + Ok(g) => g.with_num_molecules(3), + Err(e) => { + eprintln!("skipping GPU num-molecules test: {e}"); + return; + } + }; + let cpu = DirectTransform::new(pripps.formfactors.clone(), 6, None).with_num_molecules(3); + + let mut structure = Structure::from_file(&xyz, &pripps.atomkinds()).unwrap(); + structure.box_sides = Some(Vector3::new(100.0, 100.0, 100.0)); + + let cpu_i = cpu.calc_intensities(&structure).unwrap(); + let gpu_i = gpu.calc_intensities(&structure).unwrap(); + + assert_eq!(cpu_i.q, gpu_i.q, "q grids must match exactly"); + let cmp = |name: &str, a: &[f64], b: &[f64]| { + assert_eq!(a.len(), b.len(), "{name} length mismatch"); + for i in 0..a.len() { + let rel = (a[i] - b[i]).abs() / a[i].abs().max(1e-30); + assert!( + rel < 1e-3, + "{name} q={}: cpu={} gpu={} rel={rel}", + cpu_i.q[i], + a[i], + b[i] + ); + } + }; + cmp("total", &cpu_i.total, &gpu_i.total); + cmp( + "self", + cpu_i.self_term.as_ref().unwrap(), + gpu_i.self_term.as_ref().unwrap(), + ); + } } diff --git a/src/lib.rs b/src/lib.rs index 061741b..7d49ee7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -82,6 +82,18 @@ pub struct Intensity { pub total: Vec, /// Per-contribution diagnostic columns; `None` for vacuum calculations. pub contributions: Option, + /// At-concentration single-molecule form factor `P(q)`, present only when + /// `--num-molecules` (a periodic box of identical molecules) is used. For a + /// flexible molecule this is the *crowded* form factor, which an experiment + /// (dividing by a dilute `P`) cannot access. `None` otherwise. + pub form_factor: Option>, + /// Apparent (effective) structure factor `S_eff(q) = I(q)/(N·P(q))`, present + /// only with `--num-molecules`. `None` otherwise. + pub effective_structure_factor: Option>, + /// Raw self term `Σ_m |F_m(q)|² = N·P(q)` from the per-molecule partial sums; + /// the trajectory-averaging carrier from which `form_factor` / + /// `effective_structure_factor` are derived. Not part of the public output. + pub(crate) self_term: Option>, } /// Per-contribution intensity curves emitted by the three-contribution @@ -325,6 +337,7 @@ impl Pripps { xtc_opts: Option, box_sides: Option<[f64; 3]>, solvent: Option<&SolventConfig>, + num_molecules: Option, ) -> Result { let mut structure = Structure::from_file(structure_file, &self.atomkinds())?; if let Some(box_sides) = box_sides { @@ -338,6 +351,55 @@ impl Pripps { let active_solvent = solvent.filter(|cfg| cfg.is_active()); + // `--num-molecules N` (effective structure factor) needs the explicit + // periodic-box method, vacuum, and an integer N. Validate up front so a + // bad request fails immediately, before any (trajectory) work. + if let Some(n) = num_molecules { + if !matches!( + scheme, + IntensityScheme::Direct { .. } | IntensityScheme::DirectGpu { .. } + ) { + return Err(Error::usage( + "--num-molecules requires the `direct` scheme (the periodic-box method)", + )); + } + if n == 0 { + return Err(Error::usage("--num-molecules must be at least 1")); + } + // (An empty box would pass the modulus check below — `0 % n == 0` — + // and yield `S_eff = 0/0 = NaN`, but `Structure::from_file` already + // rejects atomless inputs, so `structure.pos` is non-empty here.) + if active_solvent.is_some() { + return Err(Error::usage( + "--num-molecules supports vacuum runs only (no excluded-volume/hydration model)", + )); + } + if !structure.pos.len().is_multiple_of(n) { + return Err(Error::validation(format!( + "box atom count ({}) is not a multiple of --num-molecules ({n})", + structure.pos.len() + ))); + } + // A divisible count is not enough: the N blocks must actually be + // identical contiguous molecules. Require every M-atom block to share + // the first block's atom-kind sequence — this catches a count that + // splits molecules or doesn't match the box composition (e.g. asking + // for 7 molecules from a single 1001-atom protein). + let m = structure.pos.len() / n; + if m > 0 + && !structure + .ids + .chunks(m) + .all(|block| block == &structure.ids[..m]) + { + return Err(Error::validation(format!( + "--num-molecules {n} does not partition the box into {n} identical \ + contiguous molecules: the atom composition differs between blocks" + ))); + } + log::info!("Computing structure factor for {n} molecules ({m} atoms each)"); + } + let calc: Box = match scheme { IntensityScheme::Debye { qmin, qmax, nq } => { log::info!("Using Debye formula to calculate 𝐼(𝑞) for {qmin} ≤ 𝑞 ≤ {qmax} Å⁻¹"); @@ -350,18 +412,20 @@ impl Pripps { } IntensityScheme::Direct { pmax } => { log::info!("Using direct Fourier transform with p_max = {pmax}"); - Box::new(DirectTransform::new( - self.formfactors.clone(), - pmax, - active_solvent.cloned(), - )) + let transform = + DirectTransform::new(self.formfactors.clone(), pmax, active_solvent.cloned()); + Box::new(match num_molecules { + Some(n) => transform.with_num_molecules(n), + None => transform, + }) } IntensityScheme::DirectGpu { pmax } => { #[cfg(feature = "gpu")] { // The GPU path covers the vacuum case only; with an active - // solvent or no usable adapter, fall back to the CPU Direct - // transform so results stay correct. + // solvent fall back to the CPU Direct transform so results + // stay correct. (`--num-molecules` is vacuum-only, guarded + // above, so it never collides with this branch.) if active_solvent.is_some() { log::warn!( "solvent corrections are not supported on the GPU; using CPU direct transform" @@ -372,16 +436,30 @@ impl Pripps { active_solvent.cloned(), )) } else { + // The kernel emits the per-molecule self term too, so + // `--num-molecules` runs on the GPU; fall back to the CPU + // transform (also carrying the molecule count) if no + // adapter is available. match gpu::GpuDirectTransform::new(self.formfactors.clone(), pmax) { Ok(calc) => { log::info!( "Using GPU direct Fourier transform with p_max = {pmax}" ); + let calc = match num_molecules { + Some(n) => calc.with_num_molecules(n), + None => calc, + }; Box::new(calc) as Box } Err(e) => { log::warn!("GPU unavailable ({e}); using CPU direct transform"); - Box::new(DirectTransform::new(self.formfactors.clone(), pmax, None)) + let transform = + DirectTransform::new(self.formfactors.clone(), pmax, None); + let transform = match num_molecules { + Some(n) => transform.with_num_molecules(n), + None => transform, + }; + Box::new(transform) } } } @@ -414,11 +492,28 @@ impl Pripps { } }; - if let Some(xtc_opts) = xtc_opts { - trajectory_average(calc.as_ref(), &structure, xtc_opts) + let mut result = if let Some(xtc_opts) = xtc_opts { + trajectory_average(calc.as_ref(), &structure, xtc_opts)? } else { - calc.calc_intensities(&structure) + calc.calc_intensities(&structure)? + }; + + // Derive the effective structure factor from the per-molecule self term + // (`Σ_m|F_m|² = N·P`): S_eff = I_box/(N·P) = total/self; P = self/N. The + // self term is the box self-scattering, so this is exact (no approximation). + if let (Some(n), Some(self_term)) = (num_molecules, result.self_term.take()) { + let n_f = n as f64; + result.form_factor = Some(self_term.iter().map(|&s| s / n_f).collect()); + result.effective_structure_factor = Some( + result + .total + .iter() + .zip(&self_term) + .map(|(&i, &s)| if s > 0.0 { i / s } else { f64::NAN }) + .collect(), + ); } + Ok(result) } } @@ -448,6 +543,7 @@ mod tests { None, None, Some(&solvent), + None, ) .unwrap(); @@ -477,6 +573,7 @@ mod tests { None, Some(box_sides), Some(&solvent), + None, ) .unwrap(); @@ -513,6 +610,7 @@ mod tests { None, None, Some(&scaled_solvent), + None, ) .unwrap(); let debye_model = pripps @@ -531,6 +629,7 @@ mod tests { None, Some(direct_box), Some(&scaled_solvent), + None, ) .unwrap(); let direct_model = pripps @@ -553,6 +652,7 @@ mod tests { None, None, Some(&scaled_solvent), + None, ) .unwrap(); let multipole_model = pripps @@ -593,6 +693,7 @@ mod tests { None, None, Some(&solvent), + None, ) .unwrap(); @@ -710,6 +811,7 @@ mod tests { None, None, None, + None, ) .unwrap(); @@ -765,6 +867,7 @@ mod tests { None, None, Some(&solvent), + None, ) .unwrap(); @@ -801,6 +904,7 @@ mod tests { None, None, Some(&solvent), + None, ) .unwrap(); @@ -855,6 +959,7 @@ mod tests { None, None, Some(&solvent), + None, ) .unwrap(); @@ -940,6 +1045,7 @@ mod tests { Some(opts), None, None, + None, ) .unwrap() .total @@ -980,4 +1086,200 @@ mod tests { "stride=2 must select and renormalise to frame 0" ); } + + /// `--num-molecules` structure factor: a single molecule gives S_eff ≡ 1 and + /// `p ≡ total`; a box of two identical copies has `p` (the per-molecule form + /// factor) equal to the single-molecule intensity; and the guards fire. + #[test] + fn num_molecules_structure_factor_and_guards() { + use approx::assert_relative_eq; + let dir = std::env::temp_dir(); + let yaml = dir.join("pripps_nmol_ff.yaml"); + std::fs::write( + &yaml, + "X: { formfactor: !Constant 1.0 }\nY: { formfactor: !Constant 1.0 }\n", + ) + .unwrap(); + let one = dir.join("pripps_nmol_one.xyz"); + std::fs::write(&one, "2\nd\nX 0 0 -3\nX 0 0 3\n").unwrap(); + // Two copies of the same dumbbell, well separated. + let two = dir.join("pripps_nmol_two.xyz"); + std::fs::write(&two, "4\nb\nX 0 0 -3\nX 0 0 3\nX 50 0 -3\nX 50 0 3\n").unwrap(); + // Divisible by 2 but not two identical molecules: blocks [X,Y] vs [Y,X]. + let mismatch = dir.join("pripps_nmol_mismatch.xyz"); + std::fs::write(&mismatch, "4\nm\nX 0 0 0\nY 1 0 0\nY 50 0 0\nX 51 0 0\n").unwrap(); + let pripps = Pripps::from_yaml(&yaml).unwrap(); + + let run = |file: &Path, n: usize| { + pripps + .intensity( + IntensityScheme::Direct { pmax: 30 }, + file, + None, + Some([200.0; 3]), + None, + Some(n), + ) + .unwrap() + }; + + // N = 1: I_box == self ⇒ S_eff ≡ 1, p ≡ total. + let single = run(&one, 1); + let s_eff = single.effective_structure_factor.as_ref().unwrap(); + let p = single.form_factor.as_ref().unwrap(); + for ((&se, &pi), &ti) in s_eff.iter().zip(p).zip(&single.total) { + assert_relative_eq!(se, 1.0, epsilon = 1e-9); + assert_relative_eq!(pi, ti, epsilon = 1e-9); + } + + // N = 2: the per-molecule form factor equals the single-molecule + // intensity (same q-grid, identical molecule), and S_eff is present. + let box2 = run(&two, 2); + assert_eq!(box2.form_factor.as_ref().unwrap().len(), single.total.len()); + for (&p2, &t1) in box2.form_factor.as_ref().unwrap().iter().zip(&single.total) { + assert_relative_eq!(p2, t1, epsilon = 1e-6); + } + + // Guards. + let err = |scheme, file: &Path, box_sides, solvent, n| { + pripps + .intensity(scheme, file, None, box_sides, solvent, Some(n)) + .is_err() + }; + // box atom count not a multiple of N. + assert!(err( + IntensityScheme::Direct { pmax: 2 }, + &two, + Some([200.0; 3]), + None, + 3 + )); + // non-direct scheme. + assert!(err( + IntensityScheme::Debye { + qmin: 0.0, + qmax: 0.3, + nq: 8 + }, + &two, + None, + None, + 2 + )); + // active solvent. + let solvent = SolventConfig { + excluded: ExcludedVolume::Fraser { volume_scale: 1.0 }, + ..Default::default() + }; + assert!(err( + IntensityScheme::Direct { pmax: 2 }, + &two, + Some([200.0; 3]), + Some(&solvent), + 2 + )); + // divisible count, but the blocks are not identical molecules. + assert!(err( + IntensityScheme::Direct { pmax: 2 }, + &mismatch, + Some([200.0; 3]), + None, + 2 + )); + // empty box: rejected at file load (so it never reaches the 0/0 S_eff), + // a regression guard that an atomless box can't slip through. + let empty = dir.join("pripps_nmol_empty.xyz"); + std::fs::write(&empty, "0\ne\n").unwrap(); + assert!(err( + IntensityScheme::Direct { pmax: 2 }, + &empty, + Some([200.0; 3]), + None, + 2 + )); + } + + /// `--num-molecules` over a trajectory must form S_eff as the ratio of the + /// frame-averaged box intensity to the frame-averaged self term (ratio of + /// averages), not the average of per-frame S_eff. Two distinct frames make + /// the two interpretations differ, so this pins the averaging order. + #[test] + fn num_molecules_trajectory_is_ratio_of_averages() { + use approx::assert_relative_eq; + let dir = std::env::temp_dir(); + let yaml = dir.join("pripps_nmol_traj_ff.yaml"); + std::fs::write(&yaml, "X: { formfactor: !Constant 1.0 }\n").unwrap(); + // 4 atoms = two 2-atom molecules; the template coords are overwritten + // per frame, only the atom kinds matter. + let template = dir.join("pripps_nmol_traj.xyz"); + std::fs::write(&template, "4\nt\nX 0 0 0\nX 0 0 1\nX 0 0 2\nX 0 0 3\n").unwrap(); + let pripps = Pripps::from_yaml(&yaml).unwrap(); + + // Two distinct frames (coords in Å, converted to nm for the XTC); the + // 100 Å box is constant so the per-frame q-grids align. + let frames_a: [[[f32; 3]; 4]; 2] = [ + [[0., 0., 0.], [0., 0., 4.], [30., 0., 0.], [30., 0., 4.]], + [[0., 0., 0.], [0., 0., 5.], [60., 0., 0.], [60., 0., 5.]], + ]; + let make_frame = |coords: &[[f32; 3]; 4], step: u32| { + let mut positions = Vec::with_capacity(12); + for c in coords { + positions.extend_from_slice(&[c[0] / 10.0, c[1] / 10.0, c[2] / 10.0]); + } + let mut boxvec = [0.0f32; 9]; + (boxvec[0], boxvec[4], boxvec[8]) = (10.0, 10.0, 10.0); + molly::Frame { + step, + time: step as f32, + boxvec, + precision: 1000.0, + positions, + } + }; + let frame0 = make_frame(&frames_a[0], 0); + let frame1 = make_frame(&frames_a[1], 1); + + // Per-frame reference: the raw total and self term from the static path. + let calc = DirectTransform::new(pripps.formfactors.clone(), 8, None).with_num_molecules(2); + let per_frame = |frame: &molly::Frame| { + let mut s = Structure::from_file(&template, &pripps.atomkinds()).unwrap(); + s.set_positions(frame).unwrap(); + calc.calc_intensities(&s).unwrap() + }; + let f0 = per_frame(&frame0); + let f1 = per_frame(&frame1); + let (self0, self1) = (f0.self_term.unwrap(), f1.self_term.unwrap()); + + // The full trajectory path through Pripps::intensity. + let xtc = dir.join("pripps_nmol_traj.xtc"); + let mut writer = molly::XTCWriter::create(&xtc).unwrap(); + writer.write_frame(&frame0).unwrap(); + writer.write_frame(&frame1).unwrap(); + drop(writer); + let opts = TrajectoryOptions { + xtcfile: &xtc, + stride: NonZeroUsize::new(1).unwrap(), + weights: None, + }; + let traj = pripps + .intensity( + IntensityScheme::Direct { pmax: 8 }, + &template, + Some(opts), + None, + None, + Some(2), + ) + .unwrap(); + let s_eff = traj.effective_structure_factor.unwrap(); + let p = traj.form_factor.unwrap(); + + assert_eq!(traj.q, f0.q, "static and trajectory q-grids must align"); + for i in 0..traj.q.len() { + let total_avg = 0.5 * (f0.total[i] + f1.total[i]); + let self_avg = 0.5 * (self0[i] + self1[i]); + assert_relative_eq!(s_eff[i], total_avg / self_avg, epsilon = 1e-9); + assert_relative_eq!(p[i], self_avg / 2.0, epsilon = 1e-9); + } + } } diff --git a/src/multipole/mod.rs b/src/multipole/mod.rs index e0cae85..5f6f137 100644 --- a/src/multipole/mod.rs +++ b/src/multipole/mod.rs @@ -350,6 +350,7 @@ fn vacuum_intensity( q: sampler.q_values.clone(), total, contributions: None, + ..Default::default() }) } @@ -497,6 +498,7 @@ impl MultipoleModel { excluded, hydration, }), + ..Default::default() } } diff --git a/src/trajectory.rs b/src/trajectory.rs index c4d10dc..4b443a2 100644 --- a/src/trajectory.rs +++ b/src/trajectory.rs @@ -160,6 +160,10 @@ pub fn get_side_lengths(frame: &molly::Frame) -> Option { /// `set_positions` refreshes both coordinates and (if the frame carries /// one) the periodic box on every iteration. The frame loop is serial; /// each calculator parallelizes internally over q values / q vectors. +/// +/// Curves are accumulated positionally against the q-grid of the first frame, +/// so a **fixed box** (constant volume) is assumed: a fluctuating (NPT) box +/// shifts the per-frame Direct q-grid and would misalign the average. pub fn trajectory_average( calc: &dyn IntensityCalculator, structure: &Structure, @@ -199,11 +203,19 @@ pub fn trajectory_average( if mean.q.is_empty() { mean.q = curve.q.clone(); mean.total = vec![0.0; curve.total.len()]; + // Mirror `total` for the per-molecule self term, so S_eff is formed + // as a ratio of weighted averages after the loop. + mean.self_term = curve.self_term.as_ref().map(|s| vec![0.0; s.len()]); } for (acc, v) in zip(&mut mean.total, &curve.total) { *acc += weight * v; } + if let (Some(acc), Some(v)) = (mean.self_term.as_mut(), &curve.self_term) { + for (a, x) in zip(acc, v) { + *a += weight * x; + } + } processed += 1; #[cfg(feature = "cli")] @@ -216,8 +228,22 @@ pub fn trajectory_average( 1e3 * elapsed / processed.max(1) as f64 ); - // Normalize the accumulated sum by the total weight to get the weighted average - for acc in &mut mean.total { + // The *consumed* weights can sum to zero even though `read_weights` checked + // the full set (striding may select a zero-sum subset of signed weights); + // dividing by it would turn every average into NaN, so reject it. + if !weight_sum.abs().is_normal() { + return Err(Error::validation(format!( + "weights of the {processed} averaged frames sum to {weight_sum}; cannot normalize" + ))); + } + + // Normalize the accumulated sums by the total weight to get the weighted + // averages (`total` and the self term divided by the same weight). + for acc in mean + .total + .iter_mut() + .chain(mean.self_term.iter_mut().flatten()) + { *acc /= weight_sum; }