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; }