diff --git a/README.md b/README.md index 6d56ed9..17e2594 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,17 @@ for iterative fitting against experimental data. ### Installation -Requires a Rust toolchain and Python ≥ 3.9. +Install straight from GitHub with [uv](https://docs.astral.sh/uv/). The +extension is built from source, so a +[Rust toolchain](https://www.rust-lang.org/tools/install) and Python ≥ 3.9 are +required: + +```sh +uv pip install "git+https://github.com/mlund/pripps.git#subdirectory=pripps-py" +``` + +To also get the example structures and form-factor assets used below, clone the +repo and install in editable mode instead: ```sh git clone https://github.com/mlund/pripps.git @@ -29,79 +39,90 @@ model = pr.multipole_model( qmin=0.0, qmax=0.5, nq=500, - excluded="foxs", # FoXS-style excluded volume (16π Fraser) - volume_scale=1.0, # fit knob c₁ (scales excluded volume) - hydration="sasa", # FoXS-style SASA-weighted hydration shell + excluded="foxs", # or pripps.Excluded.Foxs — 16π Fraser excluded volume + hydration="sasa", # or pripps.Hydration.Sasa probe=1.8, # SASA probe radius (Å) - contrast_density=0.03, # fit knob c₂ (scales hydration) bulk_electron_density=0.334, # e/ų, pure water ) -# Returns a dict of numpy arrays -out = model.intensity(volume_scale=1.0, contrast_density=0.03) -print(out["q"][:3], out["total"][:3]) +# Fit knobs (volume_scale, contrast_density) are applied here, not at build time. +res = model.intensity(volume_scale=1.0, contrast_density=0.03) # -> pripps.Intensity +print(res.q[:3], res.total[:3]) ``` -The returned dict contains: - -| Key | Description | -|-----|-------------| -| `q` | Scattering vector magnitudes (Å⁻¹) | -| `total` | Total intensity I(q) | -| `atoms` | Protein-only (vacuum) contribution | -| `excluded` | Excluded-volume contribution | -| `hydration` | Hydration-shell contribution | +`intensity()` returns a typed `Intensity` with NumPy-array attributes: -### Fitting against experimental data +| Attribute | Description | +|-----------|-------------| +| `.q` | Scattering vector magnitudes (Å⁻¹) | +| `.total` | Total intensity I(q) | +| `.atoms` | Protein-only (vacuum) contribution | +| `.excluded` | Excluded-volume contribution | +| `.hydration` | Hydration-shell contribution | -`multipole_model` does the expensive per-structure setup once and -caches the intermediate scattering amplitudes. Subsequent -`model.intensity(volume_scale, contrast_density)` calls only apply -the fit knobs to the cached amplitudes — typically under a -millisecond each — so wrapping them in `scipy.minimize` or a grid -search is fluent. - -For small structures or exact reference checks, `debye_model` accepts -the same solvent keyword arguments and exposes the same -`model.intensity(volume_scale, contrast_density)` interface. It caches -the exact Debye self and cross terms for atoms, excluded volume, and -hydration before fitting. -`direct_model` does the same for the explicit discrete-q transform, -caching the complex amplitudes before recombining and averaging -duplicate q magnitudes. - -A full working example is in `pripps-py/examples/fit_lysozyme.py`, -which fits lysozyme against the shipped `tests/lysozyme/lyzexp.dat` -experimental dataset and reproduces FoXS's χ²/N ≈ 0.20: +Cached models always populate `.atoms` / `.excluded` / `.hydration` (zero +where a contribution is disabled). They are `None` only when there is no +per-contribution breakdown — currently trajectory averages (the one-shot +`Pripps.intensity(..., xtc=...)`). -```sh -python pripps-py/examples/fit_lysozyme.py -``` +### Fitting against experimental data -The core pattern: +`multipole_model` does the expensive per-structure setup once and caches the +intermediate amplitudes; `model.intensity(volume_scale, contrast_density)` then +only applies the fit knobs — typically under a millisecond. There are two ways +to fit: ```python import numpy as np -from scipy.optimize import minimize # experimental data: q (Å⁻¹), I(q), σ(q) -data = np.loadtxt("experimental.dat", comments="#") -q_exp, i_exp, sigma = data[:, 0], data[:, 1], data[:, 2] +q_exp, i_exp, sigma = np.loadtxt("experimental.dat", comments="#", unpack=True) + +# (a) built-in fit — returns a typed FitResult: +fit = model.fit(q_exp, i_exp, sigma) +print(fit.volume_scale, fit.contrast_density, fit.chi2, fit.scale) +# (b) manual fit with scipy for full control: +from scipy.optimize import minimize def chi2(params): - c1, c2 = params - out = model.intensity(volume_scale=c1, contrast_density=c2) - theory = np.interp(q_exp, out["q"], out["total"]) + out = model.intensity(*params) + theory = np.interp(q_exp, out.q, out.total) scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2) return np.sum(((scale * theory - i_exp) / sigma) ** 2) / len(q_exp) +result = minimize(chi2, x0=[1.0, 0.03], + bounds=[(0.95, 1.12), (-2.0, 4.0)], method="L-BFGS-B") +``` + +`debye_model` (exact, for small structures) and `direct_model` (explicit +discrete-q) accept the same solvent arguments and expose the same +`intensity` / `fit` / `q_values` interface. -result = minimize( - chi2, - x0=[1.0, 0.2], - bounds=[(0.95, 1.12), (-2.0, 4.0)], - method="L-BFGS-B", +A full working example is in `pripps-py/examples/fit_lysozyme.py`, which fits +lysozyme against the shipped `tests/lysozyme/lyzexp.dat` and reproduces FoXS's +χ²/N ≈ 0.20 (demonstrating both fitting styles): + +```sh +python pripps-py/examples/fit_lysozyme.py +``` + +### Trajectory averaging + +To average I(q) over an MD trajectory, use the one-shot `Pripps.intensity` +(the cached models are single-structure). The `model=` argument selects the +calculator; trajectory averages carry no per-contribution breakdown: + +```python +avg = pr.intensity( + structure_file="protein.pdb", # topology / template frame + model="multipole", # or "debye" / "direct", or a pripps.Scheme + xtc="trajectory.xtc", + stride=10, # use every 10th frame + weights="weights.dat", # optional per-frame weights + excluded="foxs", + hydration="sasa", + contrast_density=0.03, ) -print(f"c1={result.x[0]:.4f}, c2={result.x[1]:.4f}, χ²/N={result.fun:.4f}") +print(avg.q, avg.total) ``` ### Solvent models @@ -110,9 +131,9 @@ print(f"c1={result.x[0]:.4f}, c2={result.x[1]:.4f}, χ²/N={result.fun:.4f}") |-----------|--------|-------------| | `excluded` | `"disabled"`, `"fraser"`, `"foxs"`, `"grid"`, `"voronoi"`, `"per-kind"` | Excluded-volume model | | `hydration` | `"disabled"`, `"sasa"`, `"grid"`, `"voronoi"` | Hydration-shell model | -| `volume_scale` | float (default 1.0) | Scales excluded volume (PepsiSAXS `r₀`) | +| `volume_scale` | float (default 1.0) | Scales excluded volume (PepsiSAXS `r₀`). In Python this is a fit knob passed to `model.intensity()` / `model.fit()`, not the model builder | | `excluded_spacing` | float (default 1.0) | Grid spacing (Å) for `excluded="grid"` | -| `contrast_density` | float (default 0.03) | Fractional hydration-layer excess density (c₂ in `A − c₁·C + c₂·B`) — pepsi's `d_rho / ρ_water`, ~0.1 for a 10 % excess layer | +| `contrast_density` | float (default 0.03) | Fractional hydration-layer excess density (c₂ in `A − c₁·C + c₂·B`) — pepsi's `d_rho / ρ_water`, ~0.1 for a 10 % excess layer. In Python a fit knob on `model.intensity()` / `model.fit()` | | `probe` | float (default: 1.8 for SASA/grid, 6.0 for voronoi) | Water probe radius in Å | | `spacing` | float (default 3.10516) | Grid spacing (Å) for `hydration="grid"`; the bulk-water linear spacing, so each cell holds ~1 water molecule | | `thickness` | float (default 3.0) | Shell thickness (Å) for `hydration="grid"` | diff --git a/assets/stovgaard2010.py b/assets/stovgaard2010.py index ef86a53..9389cb6 100644 --- a/assets/stovgaard2010.py +++ b/assets/stovgaard2010.py @@ -10,6 +10,6 @@ if residue == "q" or residue == "Unnamed: 0": continue print(f"{residue}: {{ formfactor: !Table [", end="") - for q, I in zip(df["q"], df[residue]): - print(f"[{q}, {I}], ", end="") + for q, intensity in zip(df["q"], df[residue]): + print(f"[{q}, {intensity}], ", end="") print("]}\n") diff --git a/pripps-py/examples/batch_single_bead_fit_sasbdb.py b/pripps-py/examples/batch_single_bead_fit_sasbdb.py index 9db3a88..9ee7872 100644 --- a/pripps-py/examples/batch_single_bead_fit_sasbdb.py +++ b/pripps-py/examples/batch_single_bead_fit_sasbdb.py @@ -13,8 +13,8 @@ as testing_single-bead_approximation.ipynb: - atomfile: assets/poly_amino_acid.yaml (or poly_martini.yaml for Martini3) - excluded: per-kind using the atomfile's excluded_solvent entries - - hydration: disabled - - optimizer: L-BFGS-B with bounds [(0.95, 1.12), (-2.0, 4.0)] or change them in the script + - hydration: grid + - fit: built-in `model.fit(q, I, sigma)` (pripps' library optimizer) """ from __future__ import annotations @@ -25,7 +25,6 @@ import numpy as np import pripps -from scipy.optimize import minimize def load_exp_data(path: Path, qmax: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]: @@ -54,12 +53,6 @@ def load_exp_data(path: Path, qmax: float) -> tuple[np.ndarray, np.ndarray, np.n return data[:, 0], data[:, 1], data[:, 2] -def weighted_scale(theory: np.ndarray, i_exp: np.ndarray, sigma: np.ndarray) -> float: - num = np.sum(theory * i_exp / sigma**2) - den = np.sum((theory / sigma) ** 2) - return float(num / den) - - def fit_single( pr: pripps.Pripps, xyz_path: Path, @@ -76,38 +69,25 @@ def fit_single( nq=nq, excluded="per-kind", hydration="grid", - volume_scale=1.0, probe=1.8, - contrast_density=0.03, bulk_electron_density=0.334, - spacing = 3.8, + spacing=3.8, ) - def chi2(params: np.ndarray) -> float: - c1, c2 = float(params[0]), float(params[1]) - out = model.intensity(volume_scale=c1, contrast_density=c2) - theory = np.interp(q_exp, out["q"], out["total"]) - scale = weighted_scale(theory, i_exp, sigma) - return float(np.sum(((scale * theory - i_exp) / sigma) ** 2) / len(q_exp)) - - result = minimize( - chi2, - x0=[1.0, 0.2], - bounds=[(-2.0, 3.0), (-2.0, 4.0)], - method="L-BFGS-B", - ) + # Built-in fit of (volume_scale, contrast_density) against the data. + fit = model.fit(q_exp, i_exp, sigma) - c1_opt, c2_opt = float(result.x[0]), float(result.x[1]) - out = model.intensity(volume_scale=c1_opt, contrast_density=c2_opt) - theory_interp = np.interp(q_exp, out["q"], out["total"]) - scale = weighted_scale(theory_interp, i_exp, sigma) - i_fit = scale * theory_interp + out = model.intensity( + volume_scale=fit.volume_scale, contrast_density=fit.contrast_density + ) + theory_interp = np.interp(q_exp, out.q, out.total) + i_fit = fit.scale * theory_interp return { - "c1": c1_opt, - "c2": c2_opt, - "chi2": float(result.fun), - "scale": scale, + "c1": fit.volume_scale, + "c2": fit.contrast_density, + "chi2": fit.chi2, + "scale": fit.scale, "q_exp": q_exp, "i_exp": i_exp, "sigma": sigma, diff --git a/pripps-py/examples/fit_lysozyme.py b/pripps-py/examples/fit_lysozyme.py index 94395e5..db75098 100644 --- a/pripps-py/examples/fit_lysozyme.py +++ b/pripps-py/examples/fit_lysozyme.py @@ -5,6 +5,11 @@ excluded-volume formula (16π denominator), and SASA-weighted hydration. Reproduces FoXS's χ²/N ≈ 0.20 on lysozyme. +Demonstrates BOTH fitting styles: + (a) the built-in `model.fit(...)` convenience, and + (b) a manual `scipy.optimize.minimize` over the typed `Intensity` + result (`.q` / `.total`) for full control. + Run from the repository root (after `cd pripps-py && uv pip install -e .`): python pripps-py/examples/fit_lysozyme.py @@ -32,60 +37,67 @@ # Prepare the Multipole model — expensive work runs once here. # Uses FoXS-compatible Cromer-Mann form factors with the 16π -# Fraser excluded-volume formula and SASA hydration. +# Fraser excluded-volume formula and SASA hydration. Model choice +# may be a string ("foxs") or an enum (pripps.Excluded.Foxs). pr = pripps.Pripps(atomfile="assets/foxs_formfactors.yaml") model = pr.multipole_model( structure_file="tests/lysozyme/6lyz.pdb", qmin=0.0, qmax=float(q_exp[-1]), nq=500, - excluded="foxs", - hydration="sasa", + excluded=pripps.Excluded.Foxs, + hydration=pripps.Hydration.Sasa, ) +# (a) Built-in fit — one call, returns a typed FitResult. +fit = model.fit(q_exp, i_exp, sigma) +print("Built-in model.fit():") +print(f" {fit!r}") +print(" FoXS reference: χ²/N ≈ 0.20") + +# (b) Manual fit with scipy for full control over the objective. +# The typed Intensity gives numpy `.q` / `.total` directly. def chi2(params): - """χ² between the theoretical curve and experimental data.""" c1, c2 = params out = model.intensity(volume_scale=c1, contrast_density=c2) - theory = np.interp(q_exp, out["q"], out["total"]) + theory = np.interp(q_exp, out.q, out.total) scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2) return np.sum(((scale * theory - i_exp) / sigma) ** 2) / len(q_exp) -# Grid search to find a good starting point -print("Grid search over (volume_scale, contrast_density)...") -best_chi2 = float("inf") -best_params = (1.0, 0.0) -for c1 in np.linspace(0.95, 1.12, 18): - for c2 in np.linspace(-2.0, 4.0, 61): - val = chi2([c1, c2]) - if val < best_chi2: - best_chi2 = val - best_params = (c1, c2) -print(f" Best grid point: c1={best_params[0]:.3f}, c2={best_params[1]:.3f}, χ²/N={best_chi2:.4f}") - -# Refine with L-BFGS-B result = minimize( chi2, - x0=list(best_params), + x0=[1.0, 0.0], # neutral start — independent of (a), to confirm they agree bounds=[(0.95, 1.12), (-2.0, 4.0)], method="L-BFGS-B", ) -c1_opt, c2_opt = result.x -print(f" Optimised: c1={c1_opt:.4f}, c2={c2_opt:.4f}, χ²/N={result.fun:.4f}") -print(f" FoXS reference: χ²/N ≈ 0.20") - -# Write the fitted curve -out = model.intensity(volume_scale=c1_opt, contrast_density=c2_opt) -theory = np.interp(q_exp, out["q"], out["total"]) -scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2) +print("Manual scipy.minimize():") +print(f" c1={result.x[0]:.4f}, c2={result.x[1]:.4f}, χ²/N={result.fun:.4f}") +# Write the fitted curve (using the built-in fit's parameters). +out = model.intensity( + volume_scale=fit.volume_scale, contrast_density=fit.contrast_density +) +theory = np.interp(q_exp, out.q, out.total) np.savetxt( "fit_lysozyme.csv", - np.column_stack([q_exp, i_exp, scale * theory]), + np.column_stack([q_exp, i_exp, fit.scale * theory]), delimiter=",", header="q,i_exp,i_fit", comments="", ) -print(f" Fitted curve written to fit_lysozyme.csv (scale={scale:.6e})") +print(f"Fitted curve written to fit_lysozyme.csv (scale={fit.scale:.6e})") + +# --- Trajectory averaging (one-shot; the only path that reads an XTC) --- +# Uncomment with a real trajectory to average I(q) over frames: +# avg = pr.intensity( +# structure_file="tests/lysozyme/6lyz.pdb", +# model="multipole", +# xtc="trajectory.xtc", +# stride=10, +# excluded="foxs", +# hydration="sasa", +# contrast_density=0.03, +# ) +# print(avg.q, avg.total) # note: trajectory averages have no per-contribution breakdown diff --git a/pripps-py/examples/plotting_sba_sasbdb.ipynb b/pripps-py/examples/plotting_sba_sasbdb.ipynb index 7bd7a12..249b728 100644 --- a/pripps-py/examples/plotting_sba_sasbdb.ipynb +++ b/pripps-py/examples/plotting_sba_sasbdb.ipynb @@ -112,14 +112,18 @@ " # Skip files that do not contain enough numeric data for an x/y plot.\n", " if num_df.shape[1] < 2:\n", " ax.set_title(csv_path.parent.name)\n", - " ax.text(0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\")\n", + " ax.text(\n", + " 0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\"\n", + " )\n", " ax.axis(\"off\")\n", " continue\n", "\n", " x = num_df.iloc[:, 0]\n", "\n", " # Plot each experimental intensity column, then overlay the fitted curve.\n", - " for col in num_df.columns[1:-2]: # Exclude last two columns, typically i_fit and sigma.\n", + " for col in num_df.columns[\n", + " 1:-2\n", + " ]: # Exclude last two columns, typically i_fit and sigma.\n", " ax.plot(x, num_df[col], label=col)\n", " ax.plot(x, num_df[\"i_fit\"], label=\"i_fit\", linestyle=\"--\")\n", "\n", @@ -131,7 +135,7 @@ " ax.legend(fontsize=8)\n", "\n", "# Hide unused subplot slots when the number of files is odd.\n", - "for ax in axes[len(csv_files):]:\n", + "for ax in axes[len(csv_files) :]:\n", " ax.axis(\"off\")\n", "\n", "plt.tight_layout()\n", @@ -214,18 +218,29 @@ "\n", " if num_df.shape[1] < 2:\n", " ax.set_title(csv_path.parent.name)\n", - " ax.text(0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\")\n", + " ax.text(\n", + " 0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\"\n", + " )\n", " ax.axis(\"off\")\n", " continue\n", "\n", " x = num_df.iloc[:, 0]\n", "\n", " # Draw all experimental intensity columns before adding model fits.\n", - " for col in num_df.columns[1:-2]: # Exclude last two columns, typically i_fit and sigma.\n", + " for col in num_df.columns[\n", + " 1:-2\n", + " ]: # Exclude last two columns, typically i_fit and sigma.\n", " ax.plot(x, num_df[col], label=\"Experiment\", alpha=0.5)\n", "\n", " ax.plot(x, num_df[\"i_fit\"], label=\"Amino acid bead fit\", linestyle=\"-\", color=\"m\")\n", - " ax.plot(x, num_df_martini[\"i_fit\"], label=\"Martini bead fit\", linestyle=\"-\", color=\"k\", alpha=0.5)\n", + " ax.plot(\n", + " x,\n", + " num_df_martini[\"i_fit\"],\n", + " label=\"Martini bead fit\",\n", + " linestyle=\"-\",\n", + " color=\"k\",\n", + " alpha=0.5,\n", + " )\n", "\n", " ax.set_title(csv_path.parent.name, fontsize=18)\n", " ax.set_xlabel(r\"q ($\\AA^{-1}$)\", fontsize=16)\n", @@ -236,12 +251,12 @@ " ax.tick_params(axis=\"both\", which=\"major\", labelsize=14)\n", " ax.legend(fontsize=16, loc=\"lower center\")\n", "\n", - "for ax in axes[len(csv_files):]:\n", + "for ax in axes[len(csv_files) :]:\n", " ax.axis(\"off\")\n", "\n", "# Add stable panel labels for downstream figure references.\n", "panel_labels = [\"A\", \"B\", \"C\", \"D\"]\n", - "for i, ax in enumerate(axes[:len(csv_files)]):\n", + "for i, ax in enumerate(axes[: len(csv_files)]):\n", " label = panel_labels[i] if i < len(panel_labels) else chr(ord(\"A\") + i)\n", " ax.text(\n", " 0.03,\n", @@ -332,14 +347,18 @@ "\n", " if num_df.shape[1] < 2:\n", " ax.set_title(csv_path.parent.name)\n", - " ax.text(0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\")\n", + " ax.text(\n", + " 0.5, 0.5, \"Not enough numeric columns to plot\", ha=\"center\", va=\"center\"\n", + " )\n", " ax.axis(\"off\")\n", " continue\n", "\n", " x = num_df.iloc[:, 0]\n", "\n", " # Plot experimental columns only. Fit overlays can be restored by uncommenting below.\n", - " for col in num_df.columns[1:-2]: # Exclude last two columns, typically i_fit and sigma.\n", + " for col in num_df.columns[\n", + " 1:-2\n", + " ]: # Exclude last two columns, typically i_fit and sigma.\n", " ax.plot(x, num_df[col], label=\"Experiment\", alpha=0.5)\n", " # ax.plot(x, num_df_martini[\"i_fit\"], label=\"Martini bead fit\", linestyle=\"-\", color=\"m\")\n", " # ax.plot(x, num_df[\"i_fit\"], label=\"Amino acid bead fit\", linestyle=\"-\", color=\"m\")\n", @@ -351,7 +370,7 @@ " ax.tick_params(axis=\"both\", which=\"major\", labelsize=14)\n", " ax.legend(fontsize=16)\n", "\n", - "for ax in axes[len(csv_files):]:\n", + "for ax in axes[len(csv_files) :]:\n", " ax.axis(\"off\")\n", "\n", "plt.tight_layout()\n", @@ -415,7 +434,7 @@ " # ax.set_yscale(\"log\")\n", " # ax.set_xscale(\"log\")\n", "\n", - "for ax in axes_res[len(csv_files):]:\n", + "for ax in axes_res[len(csv_files) :]:\n", " ax.axis(\"off\")\n", "\n", "plt.tight_layout()\n", diff --git a/pripps-py/examples/testing_single-bead_approximation.ipynb b/pripps-py/examples/testing_single-bead_approximation.ipynb index 1277db1..454107f 100644 --- a/pripps-py/examples/testing_single-bead_approximation.ipynb +++ b/pripps-py/examples/testing_single-bead_approximation.ipynb @@ -48,17 +48,17 @@ } ], "source": [ - "pr = pripps.Pripps(atomfile=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/assets/poly_amino_acid.yaml\")\n", + "pr = pripps.Pripps(\n", + " atomfile=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/assets/poly_amino_acid.yaml\"\n", + ")\n", "model = pr.multipole_model(\n", " structure_file=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/6lyz.xyz\",\n", " qmin=0.0,\n", " qmax=0.5,\n", " nq=196,\n", " excluded=\"per-kind\",\n", - " volume_scale=1.0, # fit knob c₁ (scales excluded volume)\n", - " hydration=\"sasa\", # FoXS-style SASA-weighted hydration shell\n", - " probe=1.8, # SASA probe radius (Å)\n", - " contrast_density=0.03, # fit knob c₂ (scales hydration)\n", + " hydration=\"sasa\", # FoXS-style SASA-weighted hydration shell\n", + " probe=1.8, # SASA probe radius (Å)\n", " bulk_electron_density=0.334, # e/ų, pure water\n", ")\n", "\n", @@ -82,16 +82,21 @@ ], "source": [ "# experimental data: q (Å⁻¹), I(q), σ(q)\n", - "data = np.loadtxt(\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/lyzexp.dat\", comments=\"#\")\n", + "data = np.loadtxt(\n", + " \"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/lyzexp.dat\",\n", + " comments=\"#\",\n", + ")\n", "q_exp, i_exp, sigma = data[:, 0], data[:, 1], data[:, 2]\n", "\n", + "\n", "def chi2(params):\n", " c1, c2 = params\n", " out = model.intensity(volume_scale=c1, contrast_density=c2)\n", - " theory = np.interp(q_exp, out[\"q\"], out[\"total\"])\n", + " theory = np.interp(q_exp, out.q, out.total)\n", " scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2)\n", " return np.sum(((scale * theory - i_exp) / sigma) ** 2) / len(q_exp)\n", "\n", + "\n", "result = minimize(\n", " chi2,\n", " x0=[1.0, 0.2],\n", @@ -120,17 +125,19 @@ ], "source": [ "out2 = model.intensity(volume_scale=result.x[0], contrast_density=result.x[1])\n", - "theory = np.interp(q_exp, out2[\"q\"], out2[\"total\"])\n", + "theory = np.interp(q_exp, out2.q, out2.total)\n", "scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2)\n", "plt.plot(q_exp, i_exp, \"o\", label=\"Experiment\")\n", - "plt.plot(out2[\"q\"], scale*out2[\"total\"], label=\"Polynomial model\")\n", + "plt.plot(out2.q, scale * out2.total, label=\"Polynomial model\")\n", "plt.legend()\n", "plt.xlabel(\"q (Å⁻¹)\")\n", "plt.ylabel(\"I(q)\")\n", "plt.yscale(\"log\")\n", - "#plt.xscale(\"log\")\n", + "# plt.xscale(\"log\")\n", "plt.title(\"SBA fit to lysozyme data\")\n", - "plt.savefig(\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/sba_fit_excl+hydration.pdf\")\n", + "plt.savefig(\n", + " \"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/sba_fit_excl+hydration.pdf\"\n", + ")\n", "plt.show()" ] }, @@ -149,23 +156,23 @@ } ], "source": [ - "pr = pripps.Pripps(atomfile=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/assets/poly_amino_acid.yaml\")\n", + "pr = pripps.Pripps(\n", + " atomfile=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/assets/poly_amino_acid.yaml\"\n", + ")\n", "model = pr.multipole_model(\n", " structure_file=\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/6lyz.xyz\",\n", " qmin=0.0,\n", " qmax=0.5,\n", " nq=196,\n", " excluded=\"fraser\", # FoXS-style excluded volume (16π Fraser)\n", - " volume_scale=1.0, # fit knob c₁ (scales excluded volume)\n", - " hydration=\"sasa\", # FoXS-style SASA-weighted hydration shell\n", - " probe=1.8, # SASA probe radius (Å)\n", - " contrast_density=0.03, # fit knob c₂ (scales hydration)\n", + " hydration=\"sasa\", # FoXS-style SASA-weighted hydration shell\n", + " probe=1.8, # SASA probe radius (Å)\n", " bulk_electron_density=0.334, # e/ų, pure water\n", ")\n", "\n", "# Returns a dict of numpy arrays\n", "out = model.intensity(volume_scale=1.0, contrast_density=0.03)\n", - "print(out[\"q\"][:3], out[\"total\"][:3])" + "print(out.q[:3], out.total[:3])" ] }, { @@ -184,16 +191,21 @@ ], "source": [ "# experimental data: q (Å⁻¹), I(q), σ(q)\n", - "data = np.loadtxt(\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/lyzexp.dat\", comments=\"#\")\n", + "data = np.loadtxt(\n", + " \"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/lyzexp.dat\",\n", + " comments=\"#\",\n", + ")\n", "q_exp, i_exp, sigma = data[:, 0], data[:, 1], data[:, 2]\n", "\n", + "\n", "def chi2(params):\n", " c1, c2 = params\n", " out = model.intensity(volume_scale=c1, contrast_density=c2)\n", - " theory = np.interp(q_exp, out[\"q\"], out[\"total\"])\n", + " theory = np.interp(q_exp, out.q, out.total)\n", " scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2)\n", " return np.sum(((scale * theory - i_exp) / sigma) ** 2) / len(q_exp)\n", "\n", + "\n", "result = minimize(\n", " chi2,\n", " x0=[1.0, 0.2],\n", @@ -222,17 +234,17 @@ ], "source": [ "out2 = model.intensity(volume_scale=result.x[0], contrast_density=result.x[1])\n", - "theory = np.interp(q_exp, out2[\"q\"], out2[\"total\"])\n", + "theory = np.interp(q_exp, out2.q, out2.total)\n", "scale = np.sum(theory * i_exp / sigma**2) / np.sum((theory / sigma) ** 2)\n", "plt.plot(q_exp, i_exp, \"o\", label=\"Experiment\")\n", - "plt.plot(out2[\"q\"], scale*out2[\"total\"], label=\"Polynomial model\")\n", + "plt.plot(out2.q, scale * out2.total, label=\"Polynomial model\")\n", "plt.legend()\n", "plt.xlabel(\"q (Å⁻¹)\")\n", "plt.ylabel(\"I(q)\")\n", "plt.yscale(\"log\")\n", - "#plt.xscale(\"log\")\n", + "# plt.xscale(\"log\")\n", "plt.title(\"SBA fit to lysozyme data\")\n", - "#plt.savefig(\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/sba_fit_no_hydration.pdf\")\n", + "# plt.savefig(\"/Users/isabelvinterbladh/Documents/HALRIC/pripps/tests/lysozyme/sba_fit_no_hydration.pdf\")\n", "plt.show()" ] }, diff --git a/pripps-py/pyproject.toml b/pripps-py/pyproject.toml index 0b33f90..6820e8b 100644 --- a/pripps-py/pyproject.toml +++ b/pripps-py/pyproject.toml @@ -7,6 +7,10 @@ name = "pripps" version = "0.1.0" description = "Python bindings for pripps (Partial Rust Implementation of PePSi)" requires-python = ">=3.9" +# numpy is a runtime dependency: every array returned by the extension +# (intensity() dict, q_values) is created through rust-numpy, which needs +# numpy's C-API loaded at call time. +dependencies = ["numpy>=1.22"] [tool.maturin] features = ["pyo3/extension-module"] diff --git a/pripps-py/src/lib.rs b/pripps-py/src/lib.rs index 7519b66..0b7a728 100644 --- a/pripps-py/src/lib.rs +++ b/pripps-py/src/lib.rs @@ -1,13 +1,268 @@ -use numpy::IntoPyArray; -use pyo3::exceptions::{PyIOError, PyRuntimeError, PyValueError}; +use numpy::{IntoPyArray, PyArray1}; +use pyo3::exceptions::{PyIOError, PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::PyDict; + +use std::num::NonZeroUsize; +use std::path::PathBuf; use ::pripps::{ - DebyeModel, DirectModel, ExcludedVolume, Hydration, Intensity, MultipoleModel, Pripps, - SolventConfig, + DebyeModel, DirectModel, ExcludedVolume, ExpData, FitResult, Hydration, Intensity, + IntensityScheme, MultipoleModel, Pripps, SolventConfig, TrajectoryOptions, fit_parameters, }; +// ---------------------------------------------------------------------------- +// Model / solvent selection enums (also accepted as plain strings everywhere) +// ---------------------------------------------------------------------------- + +/// Excluded-volume model. Pass this enum or the equivalent lowercase +/// string (e.g. `pripps.Excluded.Foxs` or `"foxs"`). +#[pyclass(name = "Excluded", module = "pripps", eq, eq_int, from_py_object)] +#[derive(Clone, Copy, PartialEq)] +enum PyExcluded { + Disabled, + Fraser, + Foxs, + Grid, + Voronoi, + PerKind, +} + +/// Hydration-shell model. Pass this enum or the equivalent lowercase +/// string (e.g. `pripps.Hydration.Sasa` or `"sasa"`). +#[pyclass(name = "Hydration", module = "pripps", eq, eq_int, from_py_object)] +#[derive(Clone, Copy, PartialEq)] +enum PyHydration { + Disabled, + Sasa, + Grid, + Voronoi, +} + +/// Scattering-intensity calculator. Pass this enum or the equivalent +/// lowercase string (e.g. `pripps.Scheme.Multipole` or `"multipole"`). +#[pyclass(name = "Scheme", module = "pripps", eq, eq_int, from_py_object)] +#[derive(Clone, Copy, PartialEq)] +enum PyScheme { + Multipole, + Debye, + Direct, +} + +/// Canonical lowercase name for an `excluded` argument that is either a +/// string, a `pripps.Excluded`, or omitted (`None` → no excluded volume). +fn excluded_name(obj: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(o) = obj else { + return Ok("none".to_string()); + }; + if let Ok(e) = o.extract::() { + return Ok(match e { + PyExcluded::Disabled => "none", + PyExcluded::Fraser => "fraser", + PyExcluded::Foxs => "foxs", + PyExcluded::Grid => "grid", + PyExcluded::Voronoi => "voronoi", + PyExcluded::PerKind => "per-kind", + } + .to_string()); + } + if let Ok(s) = o.extract::() { + return Ok(s.to_lowercase()); + } + Err(PyTypeError::new_err( + "`excluded` must be a str or pripps.Excluded", + )) +} + +/// Canonical lowercase name for a `hydration` argument that is either a +/// string, a `pripps.Hydration`, or omitted (`None` → no hydration). +fn hydration_name(obj: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(o) = obj else { + return Ok("none".to_string()); + }; + if let Ok(e) = o.extract::() { + return Ok(match e { + PyHydration::Disabled => "none", + PyHydration::Sasa => "sasa", + PyHydration::Grid => "grid", + PyHydration::Voronoi => "voronoi", + } + .to_string()); + } + if let Ok(s) = o.extract::() { + return Ok(s.to_lowercase()); + } + Err(PyTypeError::new_err( + "`hydration` must be a str or pripps.Hydration", + )) +} + +/// Resolve a `model` argument (string, `pripps.Scheme`, or omitted → +/// Multipole) to a [`PyScheme`]. +fn scheme_from(obj: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(o) = obj else { + return Ok(PyScheme::Multipole); + }; + if let Ok(e) = o.extract::() { + return Ok(e); + } + if let Ok(s) = o.extract::() { + return match s.to_lowercase().as_str() { + "multipole" => Ok(PyScheme::Multipole), + "debye" => Ok(PyScheme::Debye), + "direct" => Ok(PyScheme::Direct), + other => Err(PyValueError::new_err(format!( + "unknown model: {other} (expected `multipole`, `debye`, or `direct`)" + ))), + }; + } + Err(PyTypeError::new_err( + "`model` must be a str or pripps.Scheme", + )) +} + +// ---------------------------------------------------------------------------- +// Result types +// ---------------------------------------------------------------------------- + +/// Result of an intensity calculation: parallel NumPy arrays. +/// +/// `q` and `total` are always present. `atoms`, `excluded` and +/// `hydration` are the per-contribution curves. Cached models always +/// populate them (zero where a contribution is disabled); they are +/// `None` only when the calculation produced no per-contribution +/// breakdown — currently trajectory averages. +/// +/// The NumPy arrays are built once and shared on every attribute access +/// (a cheap reference clone), so repeated `.q` / `.total` reads in e.g. +/// plotting loops don't re-allocate. +#[pyclass(name = "Intensity", module = "pripps")] +struct PyIntensity { + q: Py>, + total: Py>, + atoms: Option>>, + excluded: Option>>, + hydration: Option>>, + n_points: usize, +} + +impl PyIntensity { + fn new(py: Python<'_>, inner: Intensity) -> Self { + let arr = |v: Vec| v.into_pyarray(py).unbind(); + let n_points = inner.q.len(); + let (atoms, excluded, hydration) = match inner.contributions { + Some(c) => ( + Some(arr(c.atoms)), + Some(arr(c.excluded)), + Some(arr(c.hydration)), + ), + None => (None, None, None), + }; + Self { + q: arr(inner.q), + total: arr(inner.total), + atoms, + excluded, + hydration, + n_points, + } + } +} + +#[pymethods] +impl PyIntensity { + /// Scattering vector magnitudes q (Å⁻¹). + #[getter] + fn q(&self, py: Python<'_>) -> Py> { + self.q.clone_ref(py) + } + + /// Total intensity I(q). + #[getter] + fn total(&self, py: Python<'_>) -> Py> { + self.total.clone_ref(py) + } + + /// Atoms-only (in-vacuo) contribution, or `None` if unavailable. + #[getter] + fn atoms(&self, py: Python<'_>) -> Option>> { + self.atoms.as_ref().map(|a| a.clone_ref(py)) + } + + /// Excluded-volume contribution, or `None` if unavailable. + #[getter] + fn excluded(&self, py: Python<'_>) -> Option>> { + self.excluded.as_ref().map(|a| a.clone_ref(py)) + } + + /// Hydration-shell contribution, or `None` if unavailable. + #[getter] + fn hydration(&self, py: Python<'_>) -> Option>> { + self.hydration.as_ref().map(|a| a.clone_ref(py)) + } + + fn __repr__(&self) -> String { + format!( + "Intensity(points={}, contributions={})", + self.n_points, + if self.atoms.is_some() { + "present" + } else { + "none" + } + ) + } +} + +/// Best-fit parameters from [`MultipoleModel.fit`] and friends. +#[pyclass(name = "FitResult", module = "pripps")] +struct PyFitResult { + inner: FitResult, +} + +#[pymethods] +impl PyFitResult { + /// Fitted excluded-volume scale (PepsiSAXS `r₀` / `c₁`). + #[getter] + fn volume_scale(&self) -> f64 { + self.inner.volume_scale + } + + /// Fitted hydration contrast density (`c₂`). + #[getter] + fn contrast_density(&self) -> f64 { + self.inner.contrast_density + } + + /// Minimised χ²/N at the best fit. + #[getter] + fn chi2(&self) -> f64 { + self.inner.chi2 + } + + /// Optimal linear scale between theory and experiment. + #[getter] + fn scale(&self) -> f64 { + self.inner.scale + } + + fn __repr__(&self) -> String { + format!( + "FitResult(volume_scale={:.4}, contrast_density={:.4}, chi2={:.4}, scale={:.4e})", + self.inner.volume_scale, self.inner.contrast_density, self.inner.chi2, self.inner.scale + ) + } +} + +// ---------------------------------------------------------------------------- +// Pripps + cached models +// ---------------------------------------------------------------------------- + +/// Scattering calculator loaded from a form-factor YAML file. +/// +/// Construct with `Pripps(atomfile)`, 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). #[pyclass(name = "Pripps", module = "pripps")] struct PyPripps(Pripps); @@ -21,9 +276,11 @@ impl PyPripps { .map_err(into_py_err) } - /// Prepare a stateful Multipole fit. Caches per-q coefficient - /// triangles so `MultipoleModel.intensity` is cheap to call - /// repeatedly with different `(volume_scale, contrast_density)`. + /// Prepare a cached Multipole model. The expensive per-structure + /// setup runs once; `MultipoleModel.intensity(volume_scale, + /// contrast_density)` and `MultipoleModel.fit(...)` are then cheap. + /// `excluded` / `hydration` accept a string or a `pripps.Excluded` / + /// `pripps.Hydration` enum (omit for none). #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( structure_file, @@ -32,14 +289,12 @@ impl PyPripps { qmax = 0.5, nq = 100, l_max = None, - excluded = "none", - volume_scale = 1.0, - excluded_spacing = 1.0, - hydration = "none", + excluded = None, + hydration = None, probe = 1.8, spacing = 3.10516, thickness = 3.0, - contrast_density = 0.03, + excluded_spacing = 1.0, bulk_electron_density = 0.334, ))] fn multipole_model( @@ -49,25 +304,21 @@ impl PyPripps { qmax: f64, nq: usize, l_max: Option, - excluded: &str, - volume_scale: f64, - excluded_spacing: f64, - hydration: &str, + excluded: Option>, + hydration: Option>, probe: f64, spacing: f64, thickness: f64, - contrast_density: f64, + excluded_spacing: f64, bulk_electron_density: f64, ) -> PyResult { - let cfg = solvent_from_kwargs( - excluded, - volume_scale, + let cfg = solvent_basis( + excluded.as_ref(), + hydration.as_ref(), excluded_spacing, - hydration, probe, spacing, thickness, - contrast_density, bulk_electron_density, )?; self.0 @@ -76,9 +327,8 @@ impl PyPripps { .map_err(into_py_err) } - /// Prepare a stateful Debye fit. Caches exact per-q Debye terms - /// so `DebyeModel.intensity` is cheap to call repeatedly with - /// different `(volume_scale, contrast_density)`. + /// Prepare a cached Debye model (exact, O(N²·Q) — best for small + /// structures). See `multipole_model` for the solvent arguments. #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( structure_file, @@ -86,14 +336,12 @@ impl PyPripps { qmin = 0.0, qmax = 0.5, nq = 100, - excluded = "none", - volume_scale = 1.0, - excluded_spacing = 1.0, - hydration = "none", + excluded = None, + hydration = None, probe = 1.8, spacing = 3.10516, thickness = 3.0, - contrast_density = 0.03, + excluded_spacing = 1.0, bulk_electron_density = 0.334, ))] fn debye_model( @@ -102,25 +350,21 @@ impl PyPripps { qmin: f64, qmax: f64, nq: usize, - excluded: &str, - volume_scale: f64, - excluded_spacing: f64, - hydration: &str, + excluded: Option>, + hydration: Option>, probe: f64, spacing: f64, thickness: f64, - contrast_density: f64, + excluded_spacing: f64, bulk_electron_density: f64, ) -> PyResult { - let cfg = solvent_from_kwargs( - excluded, - volume_scale, + let cfg = solvent_basis( + excluded.as_ref(), + hydration.as_ref(), excluded_spacing, - hydration, probe, spacing, thickness, - contrast_density, bulk_electron_density, )?; self.0 @@ -129,23 +373,21 @@ impl PyPripps { .map_err(into_py_err) } - /// Prepare a stateful direct/explicit fit. Caches per-q-vector - /// amplitudes so `DirectModel.intensity` is cheap to call - /// repeatedly with different `(volume_scale, contrast_density)`. + /// Prepare a cached direct/explicit model (discrete q-vectors in a + /// cubic box of `side_length` Å). See `multipole_model` for the + /// solvent arguments. #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( structure_file, *, pmax = 6, side_length = 200.0, - excluded = "none", - volume_scale = 1.0, - excluded_spacing = 1.0, - hydration = "none", + excluded = None, + hydration = None, probe = 1.8, spacing = 3.10516, thickness = 3.0, - contrast_density = 0.03, + excluded_spacing = 1.0, bulk_electron_density = 0.334, ))] fn direct_model( @@ -153,25 +395,21 @@ impl PyPripps { structure_file: &str, pmax: usize, side_length: f64, - excluded: &str, - volume_scale: f64, - excluded_spacing: f64, - hydration: &str, + excluded: Option>, + hydration: Option>, probe: f64, spacing: f64, thickness: f64, - contrast_density: f64, + excluded_spacing: f64, bulk_electron_density: f64, ) -> PyResult { - let cfg = solvent_from_kwargs( - excluded, - volume_scale, + let cfg = solvent_basis( + excluded.as_ref(), + hydration.as_ref(), excluded_spacing, - hydration, probe, spacing, thickness, - contrast_density, bulk_electron_density, )?; self.0 @@ -179,82 +417,207 @@ impl PyPripps { .map(PyDirectModel) .map_err(into_py_err) } -} -#[pyclass(name = "MultipoleModel", module = "pripps")] -struct PyMultipoleModel(MultipoleModel); - -#[pymethods] -impl PyMultipoleModel { - /// Compute I(q) for the given fit parameters. Returns a dict of - /// numpy arrays: `q`, `total`, `atoms`, `excluded`, `hydration`. - fn intensity<'py>( + /// One-shot intensity for a single structure or an XTC trajectory + /// average. `model` selects the calculator (`"multipole"` / + /// `"debye"` / `"direct"` or a `pripps.Scheme`). Pass `xtc=` to + /// average over a trajectory (with optional `stride` and per-frame + /// `weights`); trajectory averages carry no per-contribution + /// breakdown. `volume_scale` / `contrast_density` are baked into the + /// solvent model here (this path is not cached). + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = ( + structure_file, + *, + model = None, + xtc = None, + stride = 1, + weights = None, + box_sides = None, + excluded = None, + hydration = None, + qmin = 0.0, + qmax = 0.5, + nq = None, + l_max = None, + pmax = 6, + probe = 1.8, + spacing = 3.10516, + thickness = 3.0, + excluded_spacing = 1.0, + bulk_electron_density = 0.334, + volume_scale = 1.0, + contrast_density = 0.03, + ))] + fn intensity( &self, - py: Python<'py>, + py: Python<'_>, + structure_file: &str, + model: Option>, + xtc: Option, + stride: usize, + weights: Option, + box_sides: Option>, + excluded: Option>, + hydration: Option>, + qmin: f64, + qmax: f64, + nq: Option, + l_max: Option, + pmax: usize, + probe: f64, + spacing: f64, + thickness: f64, + excluded_spacing: f64, + bulk_electron_density: f64, volume_scale: f64, contrast_density: f64, - ) -> PyResult> { - let out = self.0.intensity(volume_scale, contrast_density); - intensity_to_dict(py, out) - } - - /// The q grid the fit was prepared over. - #[getter] - fn q_values<'py>(&self, py: Python<'py>) -> Bound<'py, numpy::PyArray1> { - self.0.q_values().to_vec().into_pyarray(py) + ) -> PyResult { + let scheme = match scheme_from(model.as_ref())? { + PyScheme::Multipole => IntensityScheme::Multipole { + qmin, + qmax, + nq, + l_max, + }, + PyScheme::Debye => IntensityScheme::Debye { + qmin, + qmax, + nq: nq.unwrap_or(100), + }, + PyScheme::Direct => IntensityScheme::Direct { pmax }, + }; + let cfg = solvent_config( + &excluded_name(excluded.as_ref())?, + volume_scale, + excluded_spacing, + &hydration_name(hydration.as_ref())?, + probe, + spacing, + thickness, + contrast_density, + bulk_electron_density, + )?; + let box_sides = match box_sides { + None => None, + Some(b) => Some(parse_box(&b)?), + }; + let stride = NonZeroUsize::new(stride) + .ok_or_else(|| PyValueError::new_err("`stride` must be >= 1"))?; + let xtc_path = xtc.map(PathBuf::from); + let weights_path = weights.map(PathBuf::from); + let xtc_opts = xtc_path.as_deref().map(|p| TrajectoryOptions { + xtcfile: p, + stride, + weights: weights_path.as_deref(), + }); + self.0 + .intensity( + scheme, + structure_file.as_ref(), + xtc_opts, + box_sides, + Some(&cfg), + ) + .map(|inner| PyIntensity::new(py, inner)) + .map_err(into_py_err) } } -#[pyclass(name = "DebyeModel", module = "pripps")] -struct PyDebyeModel(DebyeModel); +/// Macro-free helper: define a cached-model pyclass with `intensity`, +/// `fit`, and `q_values`. (Written out per type for clear docstrings.) +macro_rules! cached_model { + ($Py:ident, $Core:ty, $name:literal, $doc:literal) => { + #[doc = $doc] + #[pyclass(name = $name, module = "pripps")] + struct $Py($Core); -#[pymethods] -impl PyDebyeModel { - /// Compute I(q) for the given fit parameters. Returns a dict of - /// numpy arrays: `q`, `total`, `atoms`, `excluded`, `hydration`. - fn intensity<'py>( - &self, - py: Python<'py>, - volume_scale: f64, - contrast_density: f64, - ) -> PyResult> { - let out = self.0.intensity(volume_scale, contrast_density); - intensity_to_dict(py, out) - } + #[pymethods] + impl $Py { + /// Compute I(q) for the given fit knobs, returning an + /// `Intensity` (numpy `.q` / `.total` / `.atoms` / ...). + #[pyo3(signature = (volume_scale = 1.0, contrast_density = 0.03))] + fn intensity( + &self, + py: Python<'_>, + volume_scale: f64, + contrast_density: f64, + ) -> PyIntensity { + PyIntensity::new(py, self.0.intensity(volume_scale, contrast_density)) + } - /// The q grid the fit was prepared over. - #[getter] - fn q_values<'py>(&self, py: Python<'py>) -> Bound<'py, numpy::PyArray1> { - self.0.q_values().to_vec().into_pyarray(py) - } + /// Fit `(volume_scale, contrast_density)` to experimental + /// data. `q`/`intensity` are equal-length arrays; `sigma` + /// (uncertainties) defaults to ones (unweighted χ²). + #[pyo3(signature = (q, intensity, sigma = None))] + fn fit( + &self, + q: Vec, + intensity: Vec, + sigma: Option>, + ) -> PyResult { + run_fit(q, intensity, sigma, |c1, c2| self.0.intensity(c1, c2)) + } + + /// The q grid the model was prepared over. + #[getter] + fn q_values<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + self.0.q_values().to_vec().into_pyarray(py) + } + } + }; } -#[pyclass(name = "DirectModel", module = "pripps")] -struct PyDirectModel(DirectModel); +cached_model!( + PyMultipoleModel, + MultipoleModel, + "MultipoleModel", + "A cached Multipole model from `Pripps.multipole_model`." +); +cached_model!( + PyDebyeModel, + DebyeModel, + "DebyeModel", + "A cached Debye model from `Pripps.debye_model`." +); +cached_model!( + PyDirectModel, + DirectModel, + "DirectModel", + "A cached direct/explicit model from `Pripps.direct_model`." +); -#[pymethods] -impl PyDirectModel { - /// Compute I(q) for the given fit parameters. Returns a dict of - /// numpy arrays: `q`, `total`, `atoms`, `excluded`, `hydration`. - fn intensity<'py>( - &self, - py: Python<'py>, - volume_scale: f64, - contrast_density: f64, - ) -> PyResult> { - let out = self.0.intensity(volume_scale, contrast_density); - intensity_to_dict(py, out) - } +// ---------------------------------------------------------------------------- +// Helpers +// ---------------------------------------------------------------------------- - /// The duplicate-averaged q grid the fit was prepared over. - #[getter] - fn q_values<'py>(&self, py: Python<'py>) -> Bound<'py, numpy::PyArray1> { - self.0.q_values().to_vec().into_pyarray(py) - } +/// Build a `SolventConfig` for a cached model: the model is built at a +/// unit basis and the `volume_scale`/`contrast_density` knobs are applied +/// later in `intensity()`, so placeholder values are used here. +fn solvent_basis( + excluded: Option<&Bound<'_, PyAny>>, + hydration: Option<&Bound<'_, PyAny>>, + excluded_spacing: f64, + probe: f64, + spacing: f64, + thickness: f64, + bulk_electron_density: f64, +) -> PyResult { + solvent_config( + &excluded_name(excluded)?, + 1.0, + excluded_spacing, + &hydration_name(hydration)?, + probe, + spacing, + thickness, + 0.0, + bulk_electron_density, + ) } #[allow(clippy::too_many_arguments)] -fn solvent_from_kwargs( +fn solvent_config( excluded: &str, volume_scale: f64, excluded_spacing: f64, @@ -310,16 +673,30 @@ fn solvent_from_kwargs( }) } -fn intensity_to_dict<'py>(py: Python<'py>, intensity: Intensity) -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("q", intensity.q.into_pyarray(py))?; - dict.set_item("total", intensity.total.into_pyarray(py))?; - if let Some(c) = intensity.contributions { - dict.set_item("atoms", c.atoms.into_pyarray(py))?; - dict.set_item("excluded", c.excluded.into_pyarray(py))?; - dict.set_item("hydration", c.hydration.into_pyarray(py))?; +/// Parse a `box_sides` argument: a single float (cubic box) or a 3-tuple. +fn parse_box(o: &Bound<'_, PyAny>) -> PyResult<[f64; 3]> { + if let Ok(a) = o.extract::<[f64; 3]>() { + return Ok(a); + } + if let Ok(v) = o.extract::() { + return Ok([v; 3]); } - Ok(dict) + Err(PyValueError::new_err( + "`box_sides` must be a float (cubic) or a 3-sequence of floats", + )) +} + +/// Run the core (volume_scale, contrast_density) fit using `eval` as the +/// per-parameter intensity evaluator. +fn run_fit Intensity>( + q: Vec, + intensity: Vec, + sigma: Option>, + eval: F, +) -> PyResult { + let exp = ExpData::from_arrays(q, intensity, sigma.unwrap_or_default()).map_err(into_py_err)?; + let inner = fit_parameters(&exp, eval).map_err(into_py_err)?; + Ok(PyFitResult { inner }) } fn into_py_err(err: ::pripps::Error) -> PyErr { @@ -337,6 +714,15 @@ fn into_py_err(err: ::pripps::Error) -> PyErr { } } +/// Fast small-angle scattering (SAXS/SANS) intensities for atomic and +/// coarse-grained structures. +/// +/// Load a form-factor table with `Pripps(atomfile)`, build a cached model +/// with `Pripps.multipole_model(structure_file, ...)`, then call +/// `model.intensity(volume_scale, contrast_density)` (→ `Intensity` with +/// numpy `.q` / `.total` / `.atoms` / `.excluded` / `.hydration`) or +/// `model.fit(q, intensity, sigma)` (→ `FitResult`). For trajectory +/// averaging use the one-shot `Pripps.intensity(model=..., xtc=...)`. #[pymodule] fn pripps(m: &Bound<'_, PyModule>) -> PyResult<()> { // Bridge the core library's `log` records into Python's `logging` @@ -349,5 +735,10 @@ fn pripps(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/pripps-py/tests/test_smoke.py b/pripps-py/tests/test_smoke.py index 5c72ad8..39f1ac2 100644 --- a/pripps-py/tests/test_smoke.py +++ b/pripps-py/tests/test_smoke.py @@ -35,47 +35,79 @@ def chi2_per_point(q_theory, i_theory, exp): class TestMultipoleModel(unittest.TestCase): - def test_prepare_and_call(self): + def _model(self, **kw): pr = pripps.Pripps("../assets/poly_amino_acid.yaml") - fit = pr.multipole_model( + return pr.multipole_model( "../tests/lysozyme/4lzt.xyz", qmin=0.0, qmax=0.5, nq=50, - excluded="fraser", - hydration="sasa", + **kw, ) - out = fit.intensity(1.0, 0.03) - self.assertEqual(out["q"].shape, (50,)) - self.assertEqual(out["total"].shape, (50,)) - self.assertTrue(np.all(np.isfinite(out["total"]))) - self.assertTrue(np.all(out["atoms"] > 0)) + + def test_prepare_and_call(self): + model = self._model(excluded="fraser", hydration="sasa") + out = model.intensity(1.0, 0.03) + self.assertIsInstance(out, pripps.Intensity) + self.assertEqual(out.q.shape, (50,)) + self.assertEqual(out.total.shape, (50,)) + self.assertTrue(np.all(np.isfinite(out.total))) + self.assertIsNotNone(out.atoms) + self.assertTrue(np.all(out.atoms > 0)) + + def test_enum_selection_matches_string(self): + a = self._model( + excluded=pripps.Excluded.Fraser, hydration=pripps.Hydration.Sasa + ) + b = self._model(excluded="fraser", hydration="sasa") + self.assertTrue( + np.allclose(a.intensity(1.0, 0.03).total, b.intensity(1.0, 0.03).total) + ) + + def test_contributions_attributes_present(self): + # Cached multipole models expose the per-contribution split. + out = self._model(excluded="fraser", hydration="sasa").intensity(1.0, 0.03) + self.assertEqual(out.atoms.shape, (50,)) + self.assertEqual(out.excluded.shape, (50,)) + self.assertEqual(out.hydration.shape, (50,)) def test_refit_produces_different_curves(self): + model = self._model(excluded="fraser", hydration="sasa") + a = model.intensity(0.95, 0.02) + b = model.intensity(1.05, 0.05) + self.assertFalse(np.allclose(a.total, b.total)) + + def test_bad_excluded_raises(self): + with self.assertRaises(ValueError) as ctx: + self._model(excluded="garbage") + self.assertIn("garbage", str(ctx.exception)) + + +class TestOneShotAndTrajectory(unittest.TestCase): + def test_oneshot_intensity(self): pr = pripps.Pripps("../assets/poly_amino_acid.yaml") - fit = pr.multipole_model( + out = pr.intensity( "../tests/lysozyme/4lzt.xyz", - qmin=0.0, + model="multipole", qmax=0.5, nq=50, - excluded="fraser", + excluded=pripps.Excluded.Fraser, hydration="sasa", + contrast_density=0.03, ) - a = fit.intensity(0.95, 0.02) - b = fit.intensity(1.05, 0.05) - self.assertFalse(np.allclose(a["total"], b["total"])) + self.assertIsInstance(out, pripps.Intensity) + self.assertEqual(out.q.shape, (50,)) + self.assertTrue(np.all(np.isfinite(out.total))) - def test_bad_excluded_raises(self): + def test_bad_stride_raises(self): pr = pripps.Pripps("../assets/poly_amino_acid.yaml") - with self.assertRaises(ValueError) as ctx: - pr.multipole_model( + with self.assertRaises(ValueError): + pr.intensity( "../tests/lysozyme/4lzt.xyz", - qmin=0.0, - qmax=0.5, - nq=10, - excluded="garbage", + model="multipole", + xtc="nonexistent.xtc", + stride=0, ) - self.assertIn("garbage", str(ctx.exception)) class TestPepsiFit(unittest.TestCase): @@ -89,8 +121,7 @@ class TestPepsiFit(unittest.TestCase): χ² = 0.272 (r0=1.614, d_rho=0.0334 e/ų, probe=3 Å, thickness=3 Å, 381 shell points). - Pripps with L-BFGS-B over (c1, c2) achieves χ²/N ≈ 0.20 - (c1=1.012, c2=0.014). PepsiSAXS default d_rho is 5% of bulk + Pripps achieves χ²/N ≈ 0.20. PepsiSAXS default d_rho is 5% of bulk density = 0.05 × 0.334 ≈ 0.017 e/ų. """ @@ -112,17 +143,22 @@ def setUpClass(cls): probe=3.0, spacing=4.0, thickness=3.0, - contrast_density=cls.PEPSI_DEFAULT_CONTRAST, ) - def test_pepsi_fit_chi2(self): - """Fitted χ²/N should be competitive with PepsiSAXS (~0.27).""" + def test_builtin_fit(self): + """The built-in `model.fit` should reach a competitive χ²/N.""" + fit = self.model.fit(self.exp[:, 0], self.exp[:, 1], self.exp[:, 2]) + self.assertIsInstance(fit, pripps.FitResult) + self.assertLess(fit.chi2, 0.5, f"built-in fit χ²/N = {fit.chi2:.4f}") + + def test_manual_fit_chi2(self): + """Manual scipy fit over the typed Intensity result.""" from scipy.optimize import minimize def objective(params): c1, c2 = params out = self.model.intensity(volume_scale=c1, contrast_density=c2) - return chi2_per_point(out["q"], out["total"], self.exp) + return chi2_per_point(out.q, out.total, self.exp) result = minimize( objective, @@ -130,17 +166,17 @@ def objective(params): bounds=[(0.95, 1.12), (-2.0, 4.0)], method="L-BFGS-B", ) - chi2 = result.fun - # PepsiSAXS v3.0 gets χ²/N = 0.272; pripps gets ≈ 0.20 - self.assertLess(chi2, 0.5, f"PepsiSAXS-style fit χ²/N = {chi2:.4f}") + self.assertLess(result.fun, 0.5, f"manual fit χ²/N = {result.fun:.4f}") def test_grid_hydration_populates_contributions(self): """Grid hydration should populate atoms, excluded, and hydration.""" - out = self.model.intensity(volume_scale=1.0, contrast_density=self.PEPSI_DEFAULT_CONTRAST) - self.assertIn("atoms", out) - self.assertIn("excluded", out) - self.assertIn("hydration", out) - self.assertTrue(np.all(np.isfinite(out["total"]))) + out = self.model.intensity( + volume_scale=1.0, contrast_density=self.PEPSI_DEFAULT_CONTRAST + ) + self.assertIsNotNone(out.atoms) + self.assertIsNotNone(out.excluded) + self.assertIsNotNone(out.hydration) + self.assertTrue(np.all(np.isfinite(out.total))) if __name__ == "__main__": diff --git a/scripts/plot_formfactors.ipynb b/scripts/plot_formfactors.ipynb index 0b71616..9222990 100644 --- a/scripts/plot_formfactors.ipynb +++ b/scripts/plot_formfactors.ipynb @@ -8,7 +8,6 @@ "outputs": [], "source": [ "%matplotlib inline\n", - "import matplotlib\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import yaml" @@ -36,8 +35,8 @@ " data = yaml.load(file, Loader=yaml.FullLoader)\n", " data = data[\"H2O\"]\n", " q = [pair[0] for pair in data]\n", - " I = [pair[1] for pair in data]\n", - " plt.plot(q, I, lw=3, alpha=0.5)\n", + " intensity = [pair[1] for pair in data]\n", + " plt.plot(q, intensity, lw=3, alpha=0.5)\n", " plt.yscale(\"log\")" ] }, @@ -85,11 +84,16 @@ ], "source": [ "import mdtraj as md\n", + "\n", "traj = md.load(\"tests/4lzt/4lzt.pdb\")\n", - "sasa = md.shrake_rupley(traj, mode=\"residue\")[0] * 100.0 # -> Å^2\n", + "sasa = md.shrake_rupley(traj, mode=\"residue\")[0] * 100.0 # -> Å^2\n", "print(f\"Total SASA = {sasa.sum():.2f} Å^2 (N={len(sasa)})\")\n", "print(sasa)\n", - "np.savetxt(\"4lzt_sasa.dat\", sasa, header=\"Residue SASA for 4lzt.pdb from mdtraj.shrake_rupley()\")" + "np.savetxt(\n", + " \"4lzt_sasa.dat\",\n", + " sasa,\n", + " header=\"Residue SASA for 4lzt.pdb from mdtraj.shrake_rupley()\",\n", + ")" ] }, { diff --git a/scripts/sasbdb_bench.py b/scripts/sasbdb_bench.py index b34779e..569529f 100755 --- a/scripts/sasbdb_bench.py +++ b/scripts/sasbdb_bench.py @@ -34,10 +34,23 @@ # Additional IDs listed in Fig. S1 of the supporting information are # not yet included here — add them once Table S2 is available. SASBDB_IDS = [ - "SASDA92", "SASDAW3", "SASDEM6", "SASDF86", - "SASDJF5", "SASDJG5", "SASDJQ4", "SASDJU5", "SASDKG4", - "SASDL82", "SASDMB5", "SASDME4", "SASDMZ9", "SASDP39", - "SASDT75", "SASDT85", "SASDT95", + "SASDA92", + "SASDAW3", + "SASDEM6", + "SASDF86", + "SASDJF5", + "SASDJG5", + "SASDJQ4", + "SASDJU5", + "SASDKG4", + "SASDL82", + "SASDMB5", + "SASDME4", + "SASDMZ9", + "SASDP39", + "SASDT75", + "SASDT85", + "SASDT95", # SASDEL9 is a dummy-residue envelope model, not atomistic — # excluded from the atomistic comparison. # SASDDD3 is a hard dataset where the deposited PDB disagrees @@ -105,7 +118,9 @@ def strip_explicit_hydrogens(src, dst): # Element column is 77-78 (1-indexed); atom name is 13-16. element = line[76:78].strip() if len(line) >= 78 else "" name = line[12:16].strip() - is_h = element == "H" or (not element and name.lstrip("0123456789").startswith("H")) + is_h = element == "H" or ( + not element and name.lstrip("0123456789").startswith("H") + ) if is_h: continue fout.write(line) @@ -151,8 +166,9 @@ def preprocess(sid, refetch=False): return None, None m = fits[0]["models"][0] if m.get("type_of_model") != "atomic": - print(f"{sid}: skipping, model type = {m.get('type_of_model')}", - file=sys.stderr) + print( + f"{sid}: skipping, model type = {m.get('type_of_model')}", file=sys.stderr + ) return None, None raw_exp = d / "exp_raw.dat" @@ -179,7 +195,8 @@ def preprocess(sid, refetch=False): fixed = d / "fixed.pdb" if not fixed.is_file() or refetch: cmd = [ - "pdbfixer", str(cleaned), + "pdbfixer", + str(cleaned), "--add-atoms=heavy", "--add-residues", "--replace-nonstandard", @@ -217,8 +234,7 @@ def parse_expdata_qmax(path): return qmax -def run_pripps(exp, fixed, out_csv, log_path, pripps, atomfile, - excluded, hydration): +def run_pripps(exp, fixed, out_csv, log_path, pripps, atomfile, excluded, hydration): """Invoke pripps with fit mode. Returns (returncode, stderr_text).""" qmax = parse_expdata_qmax(exp) # Cap at 6.0 Å⁻¹ — the atomic form-factor Gaussians in the @@ -229,16 +245,23 @@ def run_pripps(exp, fixed, out_csv, log_path, pripps, atomfile, env["RUST_LOG"] = "info" cmd = [ pripps, - "-a", str(atomfile), - "-i", str(fixed), - "-f", str(exp), - "-o", str(out_csv), + "-a", + str(atomfile), + "-i", + str(fixed), + "-f", + str(exp), + "-o", + str(out_csv), "multipole", - "--qmax", f"{qmax:.3f}", + "--qmax", + f"{qmax:.3f}", # Intentionally no --nq: pripps auto-selects at the # Shannon-Nyquist rate `2·L_max + 1` (floored at 21). - "--excluded", excluded, - "--hydration", hydration, + "--excluded", + excluded, + "--hydration", + hydration, ] r = subprocess.run(cmd, capture_output=True, text=True, env=env) log_path.write_text(r.stderr) @@ -249,6 +272,7 @@ def chi2_from_fit_csv(csv_path): """Compute reduced χ² from a pripps fit CSV (q, I_exp, I_fit, σ). Pripps already applies the analytical scale factor to I_fit.""" import math + n, total = 0, 0.0 with open(csv_path) as f: for i, line in enumerate(f): @@ -273,14 +297,22 @@ def run_pepsi(pepsi_bin, fixed, exp, out_dir): is benchmarked against. Returns χ² (pepsi's own `Chi^2` line) or NaN on failure. Cached per-entry as `pepsi_fit.log`.""" import math + log_path = out_dir / "pepsi_fit.log" out_path = out_dir / "pepsi_fit.out" if not log_path.is_file(): # Pepsi writes output relative to its own cwd; run it in the entry dir. r = subprocess.run( - [str(pepsi_bin), str(fixed.resolve()), str(exp.resolve()), - "-o", str(out_path.resolve())], - cwd=str(out_dir), capture_output=True, text=True, + [ + str(pepsi_bin), + str(fixed.resolve()), + str(exp.resolve()), + "-o", + str(out_path.resolve()), + ], + cwd=str(out_dir), + capture_output=True, + text=True, ) log_path.write_text(r.stdout + "\n---STDERR---\n" + r.stderr) if r.returncode != 0: @@ -300,27 +332,37 @@ def run_pepsi(pepsi_bin, fixed, exp, out_dir): def main(): ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--only", nargs="*", help="subset of SASBDB IDs to run") - ap.add_argument("--refetch", action="store_true", - help="ignore cache and re-download everything") - ap.add_argument("--excluded", default="foxs", - choices=["disabled", "fraser", "foxs", "grid", "voronoi"]) - ap.add_argument("--hydration", default="sasa", - choices=["disabled", "sasa", "grid", "voronoi"]) ap.add_argument( - "--atomfile", type=pathlib.Path, default=FF_YAML_DEFAULT, + "--refetch", action="store_true", help="ignore cache and re-download everything" + ) + ap.add_argument( + "--excluded", + default="foxs", + choices=["disabled", "fraser", "foxs", "grid", "voronoi"], + ) + ap.add_argument( + "--hydration", default="sasa", choices=["disabled", "sasa", "grid", "voronoi"] + ) + ap.add_argument( + "--atomfile", + type=pathlib.Path, + default=FF_YAML_DEFAULT, help=f"form-factor YAML (default: {FF_YAML_DEFAULT.relative_to(REPO)})", ) ap.add_argument( - "--label", default=None, + "--label", + default=None, help="tag baked into fit output filename (default: " - "'__')", + "'__')", ) ap.add_argument( - "--with-pepsi", action="store_true", + "--with-pepsi", + action="store_true", help="also run Pepsi-SAXS on each entry for side-by-side comparison", ) ap.add_argument( - "--pepsi", default=str(pathlib.Path("~/bin/Pepsi-SAXS").expanduser()), + "--pepsi", + default=str(pathlib.Path("~/bin/Pepsi-SAXS").expanduser()), help="path to Pepsi-SAXS binary (default: ~/bin/Pepsi-SAXS)", ) args = ap.parse_args() @@ -391,8 +433,16 @@ def fmt(val, cached=False, w=8): chi2, npts = chi2_from_fit_csv(out_csv) status = "cached" else: - rc, stderr = run_pripps(exp, fixed, out_csv, log_path, pripps, - args.atomfile, args.excluded, args.hydration) + rc, stderr = run_pripps( + exp, + fixed, + out_csv, + log_path, + pripps, + args.atomfile, + args.excluded, + args.hydration, + ) if rc != 0: short = stderr.strip().splitlines()[-1] if stderr.strip() else "unknown" row = f"{sid:<9} {n_atoms:>8} {'—':>6} {'—':>8}" @@ -415,6 +465,7 @@ def fmt(val, cached=False, w=8): if results: import math + pripps_vals = sorted(r[3] for r in results if r[3] == r[3]) pepsi_vals = sorted(r[4] for r in results if r[4] == r[4]) med_p = pripps_vals[len(pripps_vals) // 2] if pripps_vals else math.nan diff --git a/src/atomkind.rs b/src/atomkind.rs index c4311c0..bca0a27 100644 --- a/src/atomkind.rs +++ b/src/atomkind.rs @@ -36,6 +36,9 @@ pub struct AtomKind { } impl AtomKind { + /// Create a kind with the given name and no physical parameters + /// (`sigma` and `displaced_volume` are both `None`). Test-only helper. + #[cfg(test)] pub fn new(name: impl Into) -> Self { Self { name: name.into(), @@ -44,12 +47,15 @@ impl AtomKind { } } + /// Human-readable name (e.g. "C", "ALA", "H2O"). pub fn name(&self) -> &str { &self.name } + /// 2 × radius (Å), or `None` if unset. pub const fn sigma(&self) -> Option { self.sigma } + /// Displaced solvent volume in ų, or `None` if unset. pub const fn displaced_volume(&self) -> Option { self.displaced_volume } @@ -63,6 +69,8 @@ impl AtomKind { // the per-atom id in the structure — skipping the borrow step // keeps the one useful bit of information and drops the rest. pub trait FindByName { + /// Return the positional index of the first kind whose name + /// matches `name`, or `None` if there is no match. fn find_name(&self, name: &str) -> Option; } diff --git a/src/cli.rs b/src/cli.rs index d7ba387..2438139 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -40,6 +40,7 @@ pub enum Commands { /// Number of q values #[clap(long, default_value = "100")] nq: usize, + /// Solvent correction parameters (excluded volume and hydration) #[clap(flatten)] solvent: SolventArgs, }, @@ -52,6 +53,7 @@ pub enum Commands { /// (the per-frame box from the XTC is used instead). #[clap(short = 'b', long = "box")] side_length: Option, + /// Solvent correction parameters (excluded volume and hydration) #[clap(flatten)] solvent: SolventArgs, }, @@ -73,6 +75,7 @@ pub enum Commands { /// auto-computed from the structure's extent and qmax. #[clap(long)] lmax: Option, + /// Solvent correction parameters (excluded volume and hydration) #[clap(flatten)] solvent: SolventArgs, }, @@ -167,22 +170,32 @@ impl SolventArgs { /// Which excluded-volume generator to use. #[derive(Copy, Clone, Debug, ValueEnum)] pub enum ExcludedKind { + /// No excluded-volume correction #[value(alias = "none")] Disabled, + /// Fraser-style per-atom Gaussian displaced solvent Fraser, + /// Fraser displaced solvent with FoXS volume parameters Foxs, + /// Grid-based displaced-solvent volume Grid, + /// Voronoi-tessellation displaced-solvent volume Voronoi, + /// Per-species (per-kind) displaced solvent volumes PerKind, } /// Which hydration-shell generator to use. #[derive(Copy, Clone, Debug, ValueEnum)] pub enum HydrationKind { + /// No hydration-shell correction #[value(alias = "none")] Disabled, + /// Solvent-accessible-surface-area (SASA) hydration shell Sasa, + /// Grid-based hydration shell Grid, + /// Voronoi-tessellation hydration shell Voronoi, } #[derive(Parser)] @@ -192,7 +205,9 @@ pub enum HydrationKind { .usage(AnsiColor::Green.on_default().bold()) .literal(AnsiColor::Cyan.on_default().bold()) .placeholder(AnsiColor::Cyan.on_default()))] +/// Top-level command line arguments. pub struct Args { + /// Scattering-intensity scheme to run #[clap(subcommand)] pub command: Commands, @@ -229,7 +244,9 @@ pub struct Args { pub verbose: bool, } -// Wrapper for main function to handle errors +/// Parse CLI arguments, run the requested intensity computation (or +/// fit), and write the output. Wraps the real `main` so errors can be +/// reported uniformly. pub fn do_main() -> Result<()> { let args = cli::Args::parse(); // SAFETY: `set_var` is unsafe because it is not thread-safe. @@ -260,6 +277,12 @@ pub fn do_main() -> Result<()> { } => { let model = pripps.debye_model(&args.input, 0.0, exp.q_max(), nq, &solvent)?; let result = crate::fit::fit_parameters(&exp, |c1, c2| model.intensity(c1, c2))?; + log::info!( + "Fitted volume_scale={:.4}, contrast_density={:.4} (χ²/N = {:.4})", + result.volume_scale, + result.contrast_density, + result.chi2 + ); let out = model.intensity(result.volume_scale, result.contrast_density); crate::fit::write_fit_csv(&args.output, &exp, &out, result.scale)?; } @@ -269,6 +292,12 @@ pub fn do_main() -> Result<()> { })?; let model = pripps.direct_model(&args.input, pmax, box_sides, &solvent)?; let result = crate::fit::fit_parameters(&exp, |c1, c2| model.intensity(c1, c2))?; + log::info!( + "Fitted volume_scale={:.4}, contrast_density={:.4} (χ²/N = {:.4})", + result.volume_scale, + result.contrast_density, + result.chi2 + ); let out = model.intensity(result.volume_scale, result.contrast_density); crate::fit::write_fit_csv(&args.output, &exp, &out, result.scale)?; } @@ -281,6 +310,12 @@ pub fn do_main() -> Result<()> { let model = pripps.multipole_model(&args.input, 0.0, exp.q_max(), nq, l_max, &solvent)?; let result = crate::fit::fit_parameters(&exp, |c1, c2| model.intensity(c1, c2))?; + log::info!( + "Fitted volume_scale={:.4}, contrast_density={:.4} (χ²/N = {:.4})", + result.volume_scale, + result.contrast_density, + result.chi2 + ); let out = model.intensity(result.volume_scale, result.contrast_density); crate::fit::write_fit_csv(&args.output, &exp, &out, result.scale)?; } diff --git a/src/error.rs b/src/error.rs index 6594fdc..069f07a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -18,7 +18,9 @@ pub enum Error { /// at every site that reads or writes a user-supplied file. #[error("{}: {source}", .path.display())] File { + /// Path of the file being read or written when the error occurred. path: PathBuf, + /// Underlying I/O error. #[source] source: std::io::Error, }, @@ -44,7 +46,14 @@ pub enum Error { /// q value outside a form factor's declared valid range. #[error("q value {q} outside form factor range [{lo}, {hi}]")] - QOutOfRange { q: f64, lo: f64, hi: f64 }, + QOutOfRange { + /// Offending q value. + q: f64, + /// Lower bound of the form factor's valid q range. + lo: f64, + /// Upper bound of the form factor's valid q range. + hi: f64, + }, /// Data-shape / invariant failures on user input: frame↔structure /// atom-count mismatch, weight/frame count mismatch, empty input diff --git a/src/fit.rs b/src/fit.rs index 5adae3b..c31ded0 100644 --- a/src/fit.rs +++ b/src/fit.rs @@ -18,8 +18,11 @@ use std::path::Path; /// Experimental SAXS data: q, I(q), σ(q). #[derive(Debug, Clone)] pub struct ExpData { + /// Scattering vector magnitudes q (Å⁻¹). pub q: Vec, + /// Measured scattering intensity I(q) at each q. pub intensity: Vec, + /// Experimental uncertainty σ(q) on each intensity (σ=1 if absent). pub sigma: Vec, } @@ -97,6 +100,38 @@ impl ExpData { }) } + /// Build from in-memory arrays. `q` and `intensity` must be the + /// same non-zero length; `sigma` is either empty (→ σ=1, unweighted + /// χ²) or the same length as `q`. + pub fn from_arrays(q: Vec, intensity: Vec, sigma: Vec) -> Result { + if q.is_empty() { + return Err(crate::Error::validation("experimental data is empty")); + } + if intensity.len() != q.len() { + return Err(crate::Error::validation(format!( + "q ({}) and intensity ({}) length mismatch", + q.len(), + intensity.len() + ))); + } + let sigma = if sigma.is_empty() { + vec![1.0; q.len()] + } else if sigma.len() == q.len() { + sigma + } else { + return Err(crate::Error::validation(format!( + "q ({}) and sigma ({}) length mismatch", + q.len(), + sigma.len() + ))); + }; + Ok(Self { + q, + intensity, + sigma, + }) + } + /// Maximum q value in the experimental data. pub fn q_max(&self) -> f64 { self.q.iter().copied().fold(0.0_f64, f64::max) @@ -202,9 +237,13 @@ pub fn chi2_per_point(theory: &Intensity, exp: &ExpData) -> f64 { /// Result of a fit against experimental data. #[derive(Debug, Clone)] pub struct FitResult { + /// Fitted excluded-volume scaling factor c₁. pub volume_scale: f64, + /// Fitted hydration-shell contrast density c₂. pub contrast_density: f64, + /// Minimised χ²/N at the fitted parameters. pub chi2: f64, + /// Optimal linear scale factor between theory and experiment. pub scale: f64, } diff --git a/src/formfactor.rs b/src/formfactor.rs index 99dceb4..2eaf897 100644 --- a/src/formfactor.rs +++ b/src/formfactor.rs @@ -41,14 +41,18 @@ struct FormFactorEntry { } impl FormFactorMap { + /// Number of atom-kind entries in the table. pub const fn len(&self) -> usize { self.entries.len() } + /// `true` if the table contains no entries. Test-only helper. + #[cfg(test)] pub const fn is_empty(&self) -> bool { self.entries.is_empty() } + /// Iterate over `(atom kind, form factor)` pairs in atom-kind id order. pub fn iter(&self) -> impl Iterator { self.entries .iter() @@ -60,10 +64,12 @@ impl FormFactorMap { &self.entries[id].kind } + /// Iterate over the atom kinds in id order. pub(crate) fn kinds(&self) -> impl Iterator { self.entries.iter().map(|entry| &entry.kind) } + /// Iterate over the form factors in id order. pub(crate) fn form_factors(&self) -> impl Iterator { self.entries.iter().map(|entry| &entry.formfactor) } @@ -79,10 +85,15 @@ impl FormFactorMap { .map(|entry| &entry.formfactor) } + /// Excluded-solvent form factor for the kind at id `id`, if one was + /// defined. Panics on out-of-range id. pub(crate) fn excluded_solvent(&self, id: usize) -> Option<&FormFactor> { self.entries[id].excluded_solvent.as_ref() } + /// Build a table from `(kind, form factor)` pairs with no excluded-solvent + /// form factors. Test-only convenience over + /// [`from_entries_with_excluded`](Self::from_entries_with_excluded). #[cfg(test)] pub(crate) fn from_entries(entries: Vec<(AtomKind, FormFactor)>) -> Self { Self::from_entries_with_excluded( @@ -93,6 +104,8 @@ impl FormFactorMap { ) } + /// Build a table from `(kind, form factor, optional excluded-solvent + /// form factor)` triples, preserving their order as the atom-kind ids. pub(crate) fn from_entries_with_excluded( entries: Vec<(AtomKind, FormFactor, Option)>, ) -> Self { @@ -143,6 +156,9 @@ impl FormFactor { Self::SingleGaussian { a, b } } + /// Inclusive `(qmin, qmax)` range over which `F(q)` is valid. `Constant` + /// and `SingleGaussian` are valid for all `q ≥ 0`; `Table` uses the first + /// and last tabulated q. pub fn qrange(&self) -> (f64, f64) { use FormFactor::*; match self { @@ -153,11 +169,6 @@ impl FormFactor { } } - pub fn in_qrange(&self, q: f64) -> bool { - let (lo, hi) = self.qrange(); - q >= lo && q <= hi - } - /// Tabulate `F(q)` at the given q values. Errors if any q is out of range. pub fn to_table(&self, qvalues: impl IntoIterator) -> Result> { qvalues @@ -166,6 +177,8 @@ impl FormFactor { .try_collect() } + /// Evaluate `F(q)` at scattering vector `q`. Errors with + /// [`Error::QOutOfRange`] if `q` is outside the [`qrange`](Self::qrange). pub fn eval(&self, q: f64) -> Result { use FormFactor::*; let (lo, hi) = self.qrange(); diff --git a/src/lib.rs b/src/lib.rs index 6ba22e6..baac710 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,8 +1,21 @@ +//! `pripps` computes small-angle X-ray scattering (SAXS) intensity `I(q)` +//! from molecular structures. Three interchangeable calculators are offered +//! via [`IntensityScheme`]: the exact Debye summation, a Direct Fourier +//! transform over a periodic box, and a Multipole (spherical-harmonic) +//! expansion. Each can apply optional solvent corrections — excluded volume +//! and hydration — configured through [`SolventConfig`]. The high-level +//! entry point is [`Pripps`], whose [`Pripps::intensity`] runs a single +//! structure or trajectory-averages over an XTC, while the `*_model` +//! constructors prepare cached, reusable fits. + // Library-only builds (no `cli` feature) have no in-crate caller for the // intensity calculators — they're only instantiated through `Pripps::intensity`. // Suppress the per-item dead-code warnings in that configuration; external // consumers (pyo3 bindings, other library users) reach them via `Pripps`. #![cfg_attr(not(feature = "cli"), allow(dead_code))] +// Force every public item to carry a doc comment; the build fails +// otherwise. Keeps the library (and its pyo3/Python surface) documented. +#![deny(missing_docs)] mod atomkind; #[cfg(feature = "cli")] @@ -10,7 +23,7 @@ pub mod cli; mod debye; mod error; mod explicit; -pub mod fit; +mod fit; mod formfactor; mod input; mod multipole; @@ -30,14 +43,18 @@ use multipole::MultipoleIntensity; use structure::Structure; use trajectory::trajectory_average; -pub use atomkind::{AtomKind, FindByName}; pub use debye::DebyeModel; pub use error::{Error, Result}; pub use explicit::DirectModel; -pub use formfactor::{FormFactor, FormFactorMap, read_formfactors, write_formfactors}; +pub use fit::{ExpData, FitResult, fit_parameters}; pub use multipole::MultipoleModel; 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}; + /// Crate-internal short alias. Not re-exported; public signatures use /// `[f64; 3]` instead so the API doesn't leak `nalgebra` types. pub(crate) type Vector3 = nalgebra::Vector3; @@ -55,8 +72,11 @@ pub(crate) type Vector3 = nalgebra::Vector3; /// diagnostic columns. It is `None` for vacuum calculations. #[derive(Debug, Clone, Default, PartialEq)] pub struct Intensity { + /// Scattering vector magnitudes (Å⁻¹). pub q: Vec, + /// Total intensity `I(q)` at each `q`. pub total: Vec, + /// Per-contribution diagnostic columns; `None` for vacuum calculations. pub contributions: Option, } @@ -64,8 +84,11 @@ pub struct Intensity { /// Multipole path, each parallel to [`Intensity::q`]. #[derive(Debug, Clone, PartialEq)] pub struct Contributions { + /// Intensity from the atoms (solute) alone. pub atoms: Vec, + /// Excluded-volume (solvent displacement) contribution. pub excluded: Vec, + /// Hydration-shell contribution. pub hydration: Vec, } @@ -75,22 +98,36 @@ pub struct Contributions { pub enum IntensityScheme { /// Debye formula sampled on `nq` points in `[qmin, qmax]` (Å⁻¹). /// O(N²·Q) — simple, exact, best for small structures. - Debye { qmin: f64, qmax: f64, nq: usize }, + Debye { + /// Lower q bound (Å⁻¹). + qmin: f64, + /// Upper q bound (Å⁻¹). + qmax: f64, + /// Number of linearly spaced q points. + nq: usize, + }, /// Direct Fourier transform on `pmax` Miller-index multiples. The /// periodic box is taken from the [`Structure`] — from a trajectory /// frame when averaging over an XTC, or via the `box_sides` override /// parameter of [`Pripps::intensity`] for single-frame calls. - Direct { pmax: usize }, + Direct { + /// 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` /// is the multipole truncation order; `None` auto-selects from /// the structure's extent and `qmax`. Multipole { + /// Lower q bound (Å⁻¹). qmin: f64, + /// Upper q bound (Å⁻¹). qmax: f64, /// `None` → auto-select at the Nyquist rate `2·l_max + 1`. nq: Option, + /// Multipole truncation order; `None` auto-selects from the + /// structure's extent and `qmax`. l_max: Option, }, } @@ -139,17 +176,17 @@ pub struct Pripps { formfactors: FormFactorMap, } +/// Options for trajectory-averaged intensity calculations over an XTC file. pub struct TrajectoryOptions<'a> { + /// Path to the XTC trajectory file. pub xtcfile: &'a path::Path, + /// Process every `stride`-th frame. pub stride: NonZeroUsize, + /// Optional path to a per-frame weights file for the average. pub weights: Option<&'a path::Path>, } impl Pripps { - pub fn formfactors(&self) -> &FormFactorMap { - &self.formfactors - } - /// Construct from a YAML form-factor definition file. pub fn from_yaml(path: &Path) -> Result { let formfactors = read_formfactors(path)?; diff --git a/src/solvent/mod.rs b/src/solvent/mod.rs index ed551dd..1876a44 100644 --- a/src/solvent/mod.rs +++ b/src/solvent/mod.rs @@ -93,7 +93,9 @@ fn water_form_factor(formfactors: &FormFactorMap, caller: &str) -> Result