Add opt-in "sapt0-bj" D3(BJ) damping set and a cached pair-term path for fitting

[Awallace3]

Summary

Adds an opt-in D3(BJ) damping set fitted to SAPT0 dispersion, plus a split of d3() that makes fitting cheap. Defaults are unchanged; no model or pipeline uses the new set.

 d3(batch, params)
   params = resolve_d3_damping_parameters(params)
-  ...CN, C6, C8/C6, distances, damping, energies (one function)
+  terms  = d3_pair_terms(batch)            # CN, C6, C8/C6, r  — damping-independent
+  return d3_pair_energies(terms, params)   # s6/s8/a1/a2 only; autograd-safe for tensor params

 resolve_d3_damping_parameters(params)
+  if params is a str: look up D3_DAMPING_PARAMETER_SETS
   None -> "sapt-pbe0-d3i" (unchanged: s6 1, s8 0.8614, a1 0.7171, a2 0.5375)

D3_DAMPING_PARAMETER_SETS = {
    "sapt-pbe0-d3i": params_intermolecular_saptpbe0_d3i,  # default, unchanged
    "sapt0-bj": params_intermolecular_sapt0_bj,           # new, opt-in
}
# s6 1, s8 1.0896, a1 0.3836, a2 2.4363 bohr
d3(batch, params="sapt0-bj")              # or d3_damping_parameters="sapt0-bj"

sapt0-bj is an MAE fit to SAPT0/aug-cc-pVDZ dispersion (Disp_aug) over a seeded 200k subsample of the 1600K train split. A disjoint 200k refit agrees to ≤ 0.018 in every parameter. The fit, validation, figures and provenance are on the qcmlforge-exp branch d3-sapt0, under analysis/d3-sapt0/fit-v1/ (results.md).

Same patch as #35 (merged into mastiff-classical), as one commit on main: 2 files, +170/−19. d3.py on main is identical to the base #35 applied to.

Evidence

Held-out D3 dispersion error vs SAPT0 (kcal/mol). Nothing below was fitted on. "MAE fit" is sapt0-bj; "MSE fit" is the same procedure under an MSE loss and is not registered.

D3 parameters	1600K test MAE (RMSE)	S66x8 MAE	S66x8 MAE @0.90	S66x8 signed @0.90	PDB13K MAE
default sapt-pbe0-d3i	0.611 (1.343)	0.694	1.603	−0.935	0.507
sapt0-bj	0.493 (1.281)	0.462	0.994	−0.425	0.388
MSE fit	0.515 (1.173)	0.529	1.117	−0.426	0.425

Sizes: the 1600K test split has 150,000 dimers, S66x8 528 points, PDB13K 13,326 dimers.

Figures are on qcmlforge-exp d3-sapt0, under analysis/d3-sapt0/fit-v1/figures/:

• fig_s66x8_disp_curves: the dispersion column of the manuscript S66x8 curves.
• fig_s66x8_disp_mae_by_scale
• fig_1600k_test_disp_parity
• fig_pdb13k_disp

PLA15 has no SAPT0 dispersion reference: the sapt0-d4(i) record's dispersion is D4(I), so it is not a check.

Speedup (Hornet, 8 CPU threads, 2,000 1600K-test dimers / 361,001 A–B pairs, two runs each):

	old d3() @ 06d1475	new d3()	d3_pair_terms once	d3_pair_energies per evaluation
ms	135.2 / 134.9	136.6 / 130.2	115.9 / 112.0	1.60 / 1.63

• d3() itself is neither faster nor slower; the default energy sum is bit-identical (−6377.2534 old and new).
• A damping-parameter evaluation on cached pair terms is 80–85× faster than calling d3(), with bit-identical energies (torch.equal asserted on 20 parameter sets).
• This is what made the 200k-dimer, 18-start, four-objective fit run in ~47 min on one node.

Tests: tests/test_dftd3_params.py (new, 8 passed):

default parameters == {s6 1, s8 0.8614, a1 0.7171, a2 0.5375}   # pinned
d3(water dimer) default == -2.912623 kcal/mol                   # value before this PR
d3(batch, "sapt0-bj") == d3(batch, resolved dict) != default    # selected only by name
d3_pair_energies(d3_pair_terms(b), p) == d3(b, p)               # bit-identical
unknown set name -> ValueError
gradients flow to s8, a1, a2

On this branch (based on main), the new tests and every existing D3 consumer pass unchanged:

• test_dftd3_params.py + test_ap3_d3_fused.py: 32 passed.
• test_ap3_fused.py, test_pt_dataset.py, test_model_io.py, test_classical_components.py, test_freeze_unfreeze.py and 6 more: 170 passed, 28 skipped.

Merge Danger

Door: two-way

Additive API plus a pure refactor of d3(). Nothing selects "sapt0-bj" unless a caller names it, and reverting removes one dict entry and re-inlines two functions.

Blast Radius: none-by-default

Every model resolves None to the unchanged default. The one behavioural widening is that resolve_d3_damping_parameters now accepts a string, so a string that used to fail with AttributeError now resolves a named set, and an unknown name raises ValueError. Adopting sapt0-bj under an already-trained -D3 model (AP3D3-FF, E5g-x + D3, CLIFF2 + D3) changes its dispersion baseline and needs a retrain; that decision is deliberately left out of this PR.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added an optional SAPT0-BJ damping parameter set for D3 energy calculations. Select it by name or continue providing parameter mappings; the SAPT-PBE0-D3I settings remain the default.
  • Pair-term calculations and damping-based energy calculations are now available as separate operations.
• Bug Fixes
  • Unknown damping parameter set names now raise a clear error.

[Awallace3]

Fit error statistics vs SAPT0 dispersion

The target is Disp_aug (SAPT0/aug-cc-pVDZ, kcal/mol) on the 1600K train split. Both sets are drawn from one seeded permutation, np.random.default_rng(20261007).permutation(N):

• fit set: part 0, 200,000 dimers / 35.8M A–B pairs.
• check set: part 1, a disjoint 200,000 dimers / 35.7M pairs. It is used only to score fits out of sample and to refit.

Errors are D3 − SAPT0 (negative = over-binding). MSE is in (kcal/mol)²; all other columns are kcal/mol.

parameters	s6	s8	a1	a2 (bohr)	fit MAE	fit MSE	fit RMSE	fit ME	fit max|err|	check MAE	check MSE	check RMSE	check ME	check max|err|
default sapt-pbe0-d3i	1	0.8614	0.7171	0.5375	0.6053	1.7692	1.3301	−0.2157	46.30	0.6081	1.7820	1.3349	−0.2140	30.88
legacy params_intermolecular_sapt0_d3i	1	0.9429	0.3399	3.0375	0.7240	2.9695	1.7232	+0.5639	52.10	0.7257	2.9962	1.7310	+0.5638	36.60
MAE fit → sapt0-bj	1	1.0896	0.3836	2.4363	0.4887	1.6267	1.2754	+0.0320	49.80	0.4906	1.6349	1.2786	+0.0325	31.49
MSE fit	1	0.9019	0.6147	1.2231	0.5106	1.3763	1.1732	+0.0147	48.83	0.5135	1.3916	1.1797	+0.0157	32.55
MAE fit, s6 free (sensitivity)	1.1105	0.8209	0.3543	2.4547	0.4874	1.6257	1.2750	+0.0247	49.58	0.4892	1.6324	1.2777	+0.0253	31.45
MSE fit, s6 free (sensitivity)	1.4310	0.0000 †	0.4902	1.2679	0.5139	1.3636	1.1677	−0.0368	47.54	0.5160	1.3751	1.1726	−0.0349	32.40

† s8 is pinned at its lower bound. With s6 free under MSE, s8 is traded away for a larger s6, which reflects a degenerate s6/s8 direction rather than information in the data. Freeing s6 improves MAE by only 0.0013 and MSE by 0.013, so s6 stays at 1.

Refit on the check set from scratch (same 18 starts):

loss	refit s8	refit a1	refit a2	max parameter shift vs fit	check score: fit params → refit params
MAE	1.0949	0.3813	2.4538	0.0175	MAE 0.49061 → 0.49060
MSE	0.9002	0.6112	1.2344	0.0112	MSE 1.39162 → 1.39156

Optimiser:

• L-BFGS-B with autograd gradients from 18 grid starts on 20k dimers, refined on the full 200k. The MAE fits get an extra Nelder–Mead polish.
• All 18 starts reached the same basin for every fit.
• No parameter sits on a bound except † above.
• Wall time on Hornet: 15 min (MSE) and 46 min (MAE, which includes the refit).

Where the error is:

• The max|err| of 46–52 kcal/mol comes from compressed opt-loose/perturb contacts with sulfonate and phosphonate anions, where SAPT0 dispersion reaches about −55 kcal/mol and D3 gives about −10 to −20.
• Under the default parameters, 1.6% of dimers have |err| > 5 kcal/mol, and the worst 0.1% carry 10% of the squared error.
• That tail is what moves the MSE fit, so the MAE fit is the one registered as sapt0-bj.

Held out, for reference (MAE; default → sapt0-bj):

• 1600K test (150,000 dimers): 0.611 → 0.493
• S66x8 (528 points): 0.694 → 0.462
• PDB13K (13,326 dimers): 0.507 → 0.388

Source: qcmlforge-exp d3-sapt0 @ fddd43bcc, analysis/d3-sapt0/fit-v1/params/fit-{mae,mse,mae-s6,mse-s6}.json. The pair terms were extracted at QCMLForge bfad81fc.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 1acb15e8-d71e-4ed5-a43a-ba6dc3c006d7

📥 Commits

Reviewing files that changed from the base of the PR and between 207ec96 and c1d3ea1.

📒 Files selected for processing (2)

• src/qcml_dftd3/d3.py
• tests/test_dftd3_params.py

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The D3 API adds an opt-in SAPT0-BJ damping parameter set and accepts registered set names. Pair-term calculation is separated from damped energy calculation. Tests cover parameter resolution, energy equivalence, and gradients.

Changes

D3 damping and energy calculation

Layer / File(s)	Summary
Named damping parameter selection src/qcml_dftd3/d3.py, tests/test_dftd3_params.py	The damping registry includes SAPT-PBE0-D3I and SAPT0-BJ. The resolver accepts registered names, returns a copy, and raises ValueError for unknown names. Tests check defaults, named selection, and copy behavior.
Pair-term and energy calculation src/qcml_dftd3/d3.py, tests/test_dftd3_params.py	d3_pair_terms returns damping-independent pair data. d3_pair_energies applies damping parameters and converts energies to kcal/mol. d3 delegates to both helpers. Tests compare helper and wrapper energies and check gradients for s8, a1, and a2.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant D3 as d3
  participant Resolver as resolve_d3_damping_parameters
  participant Terms as d3_pair_terms
  participant Energies as d3_pair_energies
  D3->>Resolver: Resolve damping parameters
  D3->>Terms: Compute pair terms
  D3->>Energies: Apply damping parameters to pair terms

Loading

Merge Risk: ⚪ Minimal · up to c1d3e

The new damping set is opt-in, and no actionable issue is established that would prevent merging after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main changes: adding the opt-in "sapt0-bj" damping set and a cached pair-term path for fitting.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}